Technical Summary: Manifold Random Features (MRFs)

Problem Statement

Random Features (RFs) are a powerful technique for approximating non-linear bi-variate functions (particularly kernels) by mapping input vectors into a lower-dimensional space where the function can be approximated via a simple dot product. While effective in Euclidean spaces, constructing efficient RFs for functions defined on general manifolds (non-Euclidean spaces) remains a significant challenge.

Standard approaches often rely on analytical derivations (e.g., trigonometric features for Gaussian kernels) which frequently fail to produce positive and bounded features. This is critical because many applications, such as linear-attention Transformers, require strictly positive kernels; standard RFs can produce negative values, leading to poor approximation quality and high variance. Furthermore, deriving analytical RF representations for general bi-variate functions on arbitrary manifolds is often mathematically intractable.

Methodology: Manifold Random Features (MRFs)

The paper proposes Manifold Random Features (MRFs), a new paradigm that leverages the discretization of manifolds and Graph Random Features (GRFs) to learn continuous fields on manifolds. The core mechanism involves a "discrete-to-continuous" learning pipeline:

1. Discretization and GRF Supervision:

   • The input manifold $\mathcal{M}$ is discretized into a finite set of points $V_N$, forming a weighted neighborhood graph $G_N$ that encodes the local geometry.
   • The target bi-variate function $F$ (e.g., a heat kernel) is approximated on this graph.
   • Graph Random Features (GRFs) are applied to the graph $G_N$. GRFs generate "signature vectors" $\phi_f(x)$ for each node $x$ by simulating ensembles of random walks. These walks update scalar loads at visited nodes based on a modulation function $f$, which is derived from the kernel's spectral coefficients.
   • These signature vectors serve as the "ground truth" or teacher signals. They are positive and bounded by construction.

2. Learning Continuous Fields:

   • A neural network $g_{\theta}^{\mathcal{M},f}(x, \omega)$ is trained to regress the discrete signature vectors. The input to the network is a pair of points $(x, \omega)$ on the manifold, and the output approximates the value of the signature vector $\phi_f(x)$ at node $\omega$.
   • The network is trained to minimize the loss between its output and the GRF signature values. Crucially, the output is clamped at zero to ensure positivity.

3. Inference and Feature Construction:

   • Once trained, $g_{\theta}^{\mathcal{M},f}$ acts as a smooth surrogate for the discrete signatures.
   • For any new point $z \in \mathcal{M}$, the MRF vector $\phi_{\mathcal{M},F}(z)$ is constructed by evaluating $g_{\theta}^{\mathcal{M},f}(z, \omega)$ at a set of sampled anchor points $\{\omega_i\}$.
   • The approximation is given by:
     $$F(x, y) \approx \sum_{i=1}^m g_{\theta}^{\mathcal{M},f}(x, \omega_i) g_{\theta}^{\mathcal{M},f}(y, \omega_i)$$
   • This mechanism provides positive and bounded features, a key property for low-variance approximation.

Key Contributions

1. A New Paradigm for Manifold Kernels: The paper introduces MRFs, a method to approximate bi-variate functions on general manifolds by learning continuous fields via discrete GRF supervision. This bypasses the need for analytical derivations that are often impossible for complex manifolds.
2. Positive and Bounded Features: Unlike standard RFs (e.g., Random Fourier Features) which can produce negative values, MRFs inherently provide positive and bounded features. This is achieved by using GRFs as teachers and enforcing non-negativity during inference.
3. Asymptotic Connection: The authors establish a deep asymptotic connection between discrete GRFs (defined on graph objects) and continuous random features. They demonstrate that as the grid resolution approaches infinity, the discrete signature vectors converge to continuous feature maps.
4. Rediscovery of Gaussian Kernel Approximation: As a by-product, the method applied to Euclidean grid graphs re-discovers a recently proposed mechanism for approximating Gaussian kernels with positive and bounded features. This offers a simplified derivation via random walks on graphs, bypassing complex original mathematical computations.
5. Rigorous Theory and Empirical Validation: The work is supported by theoretical analysis regarding the convergence of discrete Laplacians to continuous operators and extensive experimental validation.

Experimental Results

The authors validate MRFs across several scenarios:

• Euclidean Limit (Sanity Check): On regular Euclidean grids, MRFs successfully converge to the analytic positive Gaussian feature map as the grid is refined, verifying the theoretical discrete-to-continuous transition.
• 2D Surfaces in 3D Space: MRFs were tested on embedded non-linear geometries (Sphere, Ellipsoid, Möbius strip, Torus).
  • Accuracy: MRFs significantly outperformed ambient Euclidean RF baselines (RFF, ORF, PRF, PORF) in reconstructing intrinsic manifold heat kernels, even when baselines were allowed up to 4$\times$ more features.
  • Speed: MRFs achieved inference speedups of 37x to 61x compared to brute-force spectral decomposition methods.
• Mesh Interpolation:
  • Vertex Normal Prediction: On the Thingi10k dataset, MRFs achieved higher cosine similarity than other RF methods while maintaining stable preprocessing times as mesh size increased. In contrast, full-kernel construction costs grew rapidly ($O(|V|^2)$ memory), becoming impractical for large meshes.
  • Velocity Prediction: On the flag_simple dataset, MRFs showed superior scaling exponents and inference speed (exceeding 10$\times$ speedup) compared to dense spectral baselines.
• Attention on Manifold-Valued Tokens: In a masked reconstruction task using tokens sampled from a Möbius strip, MRFs achieved the lowest Mean Squared Error (MSE) and relative $\ell_2$ error, outperforming both exact kernel approximations and other RF methods. This suggests MRFs are effective for linearized attention when token embeddings lie on non-Euclidean geometries.
• High-Dimensional Descriptors: On a Brodatz texture classification task using 15-dimensional log-SPD covariance descriptors, MRFs achieved a mean group accuracy of 55.7%, significantly outperforming the second-best method (RFF at 43.5%).

Significance and Claims

The paper claims that MRFs provide a scalable, accurate, and theoretically grounded framework for handling bi-variate functions on manifolds. The primary significance lies in:

• Geometry Awareness: MRFs respect the intrinsic geometry of the manifold, leading to superior performance compared to methods that ignore the manifold structure (ambient Euclidean RFs).
• Scalability: By avoiding explicit construction of large kernel matrices (which requires cubic time for spectral decomposition and quadratic memory), MRFs enable sub-quadratic time complexity for kernel evaluations and downstream tasks.
• Robustness: The provision of positive and bounded features addresses a critical limitation of existing RF methods, particularly for applications like linear attention where negative values degrade performance.
• Generalizability: The framework is applicable to a wide range of manifolds, including compact surfaces (spheres, tori) and, with local-on-compacts exhaustion, non-compact manifolds (demonstrated on hyperbolic 3-space).

The authors note that MRFs require an offline preprocessing stage (graph construction, random walk supervision, and neural network training). They argue this cost is justified for scenarios involving many out-of-sample evaluations or downstream inference steps, though it may not be optimal for one-off small problems. The work is presented as a methodological framework, with future work suggested for extending the theory to broader classes of non-compact manifolds.