Shape correspondence stops paying for mesh resolution: recasting spatial-spectral coupling as attention cuts peak training memory from 106.08 to 5.67 GiB
Synopsis
Dense shape correspondence on high-resolution 3D meshes has been held back by the quadratic cost of spatial-spectral coupling: for two shapes with k vertices each, a single coupling occupies roughly GB in single precision, so existing methods downsample heavily and transfer predictions back. Geo-FlashAttention rewrites that coupling as scaled dot-product attention and evaluates the exact softmax only over surface patches selected by geodesic distance, giving near-linear empirical scaling while preserving the correspondence accuracy of exact implementations. At vertices, peak training memory falls from 106.08 to 5.67 GiB, and each training step is faster than FlashAttention, making end-to-end correspondence learning feasible on meshes such as cortical surfaces.
Interpretation
Dense shape correspondence can now be trained end to end on meshes with tens of thousands to more than a hundred thousand vertices, instead of being computed on heavily downsampled meshes and transferred back. Earlier functional-map methods coupled the spectral representation back to the spatial domain by explicitly building a dense similarity matrix between the vertices of the two shapes; both forward and backward passes materialized it, so memory grew quadratically with mesh size. At vertices, peak training memory drops from 106.08 GiB to 5.67 GiB, and Geo-FlashAttention is faster per training step than FlashAttention.
The approach treats spatial-spectral coupling directly as one attention computation, then uses geometry to restrict attention to local surface patches. Scaled dot-product attention had been used for sequence modeling, where exact implementations rely on hardware-aware tiling to avoid materializing the attention matrix, while approximate methods such as linear attention replace the softmax with a low-rank feature map. Learned shape descriptors act as queries and keys and spectral basis functions as values, so the coupling is equivalent to scaled dot-product attention; Geo-FlashAttention partitions the surface with geodesic farthest-point sampling, and each target patch computes the exact softmax over only its highest-scoring source patches.
Across six standard benchmarks and three functional-map pipelines, the efficient couplings preserve the correspondence accuracy of exact implementations. Linearized couplings reduce cost to linear but change the correspondence itself, with clearly larger errors on strongly non-isometric shape pairs. On FAUST, SCAPE, SHREC'19, SMAL, DeformingThings4D and TOPKIDS, mean-error differences between dense and alternative couplings vary in direction and are generally small relative to seed variability; linear attention is significantly worse on four of the six benchmarks and Performer on five, with the largest gaps on SMAL and DT4D-inter.
Perspective
The demonstrated scope covers the ULRSSM, HFM and ZoomOut pipelines on six standard benchmarks (FAUST, SCAPE, SHREC'19, SMAL, DeformingThings4D and TOPKIDS), plus cortical matching on right-hemisphere white surfaces of Mindboggle-101 with more than a hundred thousand vertices; timing and memory were measured on a single, exclusively allocated H200 NVL.
The reformulation holds only when correspondence is consumed through a matrix product; Geo-FlashAttention depends on a fixed candidate budget and provides no approximation guarantees, and its scaling is measured empirically. Cortical evaluation lacks dense ground truth and relies on the manually corrected DKT31 parcellation and a label-free area-distortion score, while the reference labels were FreeSurfer-initialized, so agreement with that baseline is partly built into the metric.