Function svd_pe

python/dgl/transforms/functional.py:4012–4086

r"""SVD-based Positional Encoding, as introduced in `Global Self-Attention as a Replacement for Graph Convolution `__ This function computes the largest :math:`k` singular values and corresponding left and right singular vectors to form position

(g, k, padding=False, random_flip=True)

def svd_pe(g, k, padding=False, random_flip=True):
    r"""SVD-based Positional Encoding, as introduced in
    `Global Self-Attention as a Replacement for Graph Convolution
    <https://arxiv.org/pdf/2108.03348.pdf>`__

    This function computes the largest :math:`k` singular values and
    corresponding left and right singular vectors to form positional encodings.

    Parameters
    ----------
    g : DGLGraph
        A DGLGraph to be encoded, which must be a homogeneous one.
    k : int
        Number of largest singular values and corresponding singular vectors
        used for positional encoding.
    padding : bool, optional
        If False, raise an error when :math:`k > N`,
        where :math:`N` is the number of nodes in :attr:`g`.
        If True, add zero paddings in the end of encoding vectors when
        :math:`k > N`.
        Default : False.
    random_flip : bool, optional
        If True, randomly flip the signs of encoding vectors.
        Proposed to be activated during training for better generalization.
        Default : True.

    Returns
    -------
    Tensor
        Return SVD-based positional encodings of shape :math:`(N, 2k)`.

    Example
    -------
    >>> import dgl

    >>> g = dgl.graph(([0,1,2,3,4,2,3,1,4,0], [2,3,1,4,0,0,1,2,3,4]))
    >>> dgl.svd_pe(g, k=2, padding=False, random_flip=True)
    tensor([[-6.3246e-01, -1.1373e-07, -6.3246e-01,  0.0000e+00],
            [-6.3246e-01,  7.6512e-01, -6.3246e-01, -7.6512e-01],
            [ 6.3246e-01,  4.7287e-01,  6.3246e-01, -4.7287e-01],
            [-6.3246e-01, -7.6512e-01, -6.3246e-01,  7.6512e-01],
            [ 6.3246e-01, -4.7287e-01,  6.3246e-01,  4.7287e-01]])
    """
    n = g.num_nodes()
    if not padding and n < k:
        raise ValueError(
            "The number of singular values k must be no greater than the "
            "number of nodes n, but " + f"got {k} and {n} respectively."
        )
    a = g.adj_external(ctx=g.device, scipy_fmt="coo").toarray()
    u, d, vh = scipy.linalg.svd(a)
    v = vh.transpose()
    m = min(n, k)
    topm_u = u[:, 0:m]
    topm_v = v[:, 0:m]
    topm_sqrt_d = sparse.diags(np.sqrt(d[0:m]))
    encoding = np.concatenate(
        ((topm_u @ topm_sqrt_d), (topm_v @ topm_sqrt_d)), axis=1