PaperlayerView exact revision on arXiv

Source-linked AI summary

Functional Attention: From Pairwise Affinities to Functional Correspondences

Jiefang Xiao, Maolin Gao, Simon Weber, Guandao Yang, Daniel Cremers

arXiv:2605.31559v1cs.LG

TL;DR

Operator-learning models must map continuous function spaces without relying solely on token-wise representations that can miss global structure. The paper introduces Functional Attention, which uses adaptive bases and structured linear operators to represent functional correspondences. Across PDE solving, 3D segmentation, and regression, it reports state-of-the-art performance and robustness across datasets and resolutions.

  • Problem

    Operator learning addresses mappings between infinite-dimensional function spaces, but practical attention models process finite samples and may not learn operators directly.

  • Method

    Functional Attention reinterprets attention as a functional correspondence between learned function spaces, using adaptive bases and structured linear operators instead of pointwise affinities.

  • Results

    Across PDE solving, 3D segmentation, and regression, FUNCATTN consistently achieves state-of-the-art performance and strong robustness across datasets and resolutions.

  • Takeaways & Limitations

    Functional Attention provides a compact transport mechanism that is continuous, basis-aware, and naturally stable under regularization.

  • Takeaways & Limitations

    The learned basis uses a simple softmax projection, while approximation guarantees, generalization bounds, and formal links between compression ratio k/n and approximation error remain open.

Abstract

from arXiv · show

Learning mappings between infinite-dimensional function spaces, or operator learning, is essential for many machine learning applications. Although transformer-based operators are popular, they often rely on token-wise attention. These methods treat continuous fields as discrete tokens and usually ignore the global functional structure. We introduce \emph{Functional Attention}, which reinterprets attention as a functional correspondence between adaptive bases. Inspired by geometric functional maps, our method replaces softmax affinities with structured linear operators. This yields a compact, generalizable, resolution-invariant representation that explicitly captures global dependencies. Experiments demonstrate that \emph{Functional Attention} can match state-of-the-art performance in many operator learning tasks, including solving PDEs, 3D segmentation, and regression, while remaining robust to varying discretizations. Project page is available at https://github.com/xjffff/FUNCATTN.

1. Introduction

Operator learning targets mappings between infinite-dimensional function spaces, but token-wise attention can scale poorly and overlook global functional structure. Functional Attention reframes attention as a correspondence between learned function spaces and achieves strong performance and robustness across operator-learning tasks.

  • Operator learning maps infinite-dimensional function spaces and supports applications including PDE solving, computational design, and physical simulation.
  • Token-based attention represents functions through discrete samples, which may scale poorly, ignore global functional structure, and weaken consistency across resolutions or irregular meshes.
  • Functional Attention reinterprets attention as a functional correspondence between learned function spaces, inspired by the functional maps framework.
  • FUNCATTN provides an attention formalism designed to capture intrinsic structural properties and offers a unified, theoretically grounded approach to operator learning.
  • Across PDE solving, 3D segmentation, and regression, FUNCATTN consistently achieves state-of-the-art performance and exhibits robustness across datasets and resolutions.

2. Related Work

Related work addresses operator learning through spectral, graph-based, transformer, linear, and low-rank attention mechanisms. Functional Attention differs by incorporating functional structure into attention and empirically demonstrating its effectiveness.

  • Neural operator methods use spectral, graph-based, and transformer architectures to model mappings between infinite-dimensional function spaces while balancing efficiency and geometric flexibility.
  • Scaled dot-product attention has quadratic complexity in context length, whereas linear attention reorders kernelized computations to reduce complexity from quadratic to linear.
  • Functional Attention incorporates functional structures into attention and is demonstrated to be more effective empirically than the compared linear-attention perspective.
  • Low-rank attention methods reduce computational cost through projections, softmax approximations, fixed queries, hashing, or structured matrix constraints.
  • Compared with slice-and-attend methods, FUNCATTN uses a more general spectral framework whose learned functional coefficients capture intrinsic structures beyond pure physics.

3. The Problem & Motivation

Operator learning seeks discretization-independent mappings between continuous functions, while practical attention models process finite samples as tokens and can miss their geometric relationships. FUNCATTN addresses this by replacing explicit pointwise affinities with a compact, basis-aware linear operator between learned functional spaces.

  • 3.1. Operator Learning Formulation: Operator learning formalizes the target as a nonlinear operator O: F → G mapping input functions to output functions.
  • 3.1. Operator Learning Formulation: Formulating learning as an operator makes the target independent of a particular discretization, supporting resolution-invariant generalization across meshes or sampling densities.
  • 3.2. Tokenized Representations: Practical attention architectures evaluate functions at finite locations, stack samples into token matrices, and compute queries, keys, and values with learnable projections.
  • 3.2. Tokenized Representations: Token-wise attention tightly couples complexity to token count and treats dense pointwise affinities as the representation of correspondences.
  • 3.3. Motivation: The paper asks whether attention can represent the induced linear operator between function spaces directly rather than using explicit point-to-point correspondences.
  • 3.3. Motivation: Functional maps motivate representing correspondences between spaces through linear operators acting on function spaces, reducing a combinatorial problem to convex optimization.
  • 3.3. Motivation: FUNCATTN estimates a compact linear operator between learned functional spaces, replacing softmax pointwise affinities with a basis-aware operator learned through least-squares objectives.

4. Functional Attention

Functional Attention replaces token-level affinity estimation with compact, basis-aware transport between learned function spaces. Its regularized operator and adaptive bases support continuity, stability, and structured generalization.

  • Functional correspondence: FUNCATTN represents correspondence with a compact k × k linear operator instead of an n × n affinity matrix.The operator acts on functional bases associated with query and key-value spaces.
  • Functional correspondence: The method estimates the transport operator through Tikhonov-regularized least squares and applies it to spectral value coefficients.Regularization strength λ controls numerical stabilization and the resulting operator behavior.
  • Properties: The compact operator imposes an implicit low-rank constraint that can improve generalization on structured data.The framework is described as continuous, basis-aware, and naturally stable under regularization.
  • Adaptive basis: FUNCATTN uses learned data-dependent bases rather than fixed spectral bases, allowing the partition geometry to adapt to input features.The learned bases are produced from input features using a linear layer followed by softmax across the basis dimension.
  • Adaptive basis: Softmax basis functions remain bounded and sum to one, while their low-temperature limit recovers classical piecewise-constant elements.This connects the learned basis construction to generalized P0 finite-element-like functions.
  • Empirical behavior: In few-shot sinusoidal regression, FUNCATTN achieves the lowest MSE across varying context sizes and scales most effectively as context grows.The figure compares predictions at context length 4 and MSE across different context sizes.
  • Properties: For any λ > 0, the functional attention layer has a finite local Lipschitz bound whose upper bound is controlled by the regularization parameter.The bound also depends polynomially on input size, input norm, and relevant weight norms.

5. Experiments

FUNCATTN is evaluated across regression, PDE solving, 3D segmentation, OOD generalization, super-resolution, and basis ablations. It consistently performs strongly, including state-of-the-art results across diverse physical systems and improved generalization under sparse observations, unseen regimes, and finer resolutions.

  • Few-shot regression: Up to three orders of magnitude lower error than vanilla attention and Transolver, and one order lower than Intention, is achieved in few-shot sinusoidal regression.FUNCATTN reaches lower error with five observations than scaled dot-product attention with forty observations.
  • PDE solving: Five of six PDE benchmarks receive the best performance, with relative improvements of 6%–26.3% over Transolver.The benchmarks span fluid and solid mechanics across point clouds, structured meshes, and unstructured meshes.
  • 3D segmentation: The highest accuracy on RNA point-cloud segmentation surpasses PointNet++, DiffusionNet, and Transolver.The dataset contains 640 ribosomal RNA structures, each represented by 4096 points with 259 functional categories.
  • OOD generalization: 23.4 relative error and 99.4% Spearman correlation on OOD Reynolds outperform the closest competitor by 8.8%, while OOD Angles reaches 13.3% error and 99.7% correlation with a 9.5% improvement.The test sets contain unseen Reynolds numbers and angles of attack.
  • Complex geometry: 0.64% relative L2 error on Darcy flow with a triangular notch improves 30.9% over WNO, while dgFNO+ records 7.82%.The notch tip creates sharp local solution features that challenge fixed-basis spectral methods.
  • Resolution and ablations: Training at 2048 grid points and testing at 8192 without fine-tuning preserves strong performance, while freely learned bases outperform fixed Fourier and orthogonally constrained alternatives.The number of bases generally helps until potential overfitting causes slight degradation; k = 64 is recommended as a robust default within 5% of the best across benchmarks.

6. Conclusion

FUNCATTN bridges functional map theory and attention to capture functional structure in operator learning. The method is theoretically analyzed and performs strongly across PDE, segmentation, and domain-shift settings, while several theoretical and design questions remain open.

  • FUNCATTN reframes attention as a functional operator through a principled connection between functional map theory and attention mechanisms.
  • The method provides theoretical Lipschitz continuity with respect to input functions, establishing stability.
  • FUNCATTN achieves state-of-the-art accuracy on PDE benchmarks and superior generalization under domain shifts.
  • On complex geometries with sharp local features, FUNCATTN adapts to nontrivial domains, while 3D segmentation results exceed competing approaches.
  • The learned basis currently uses simple softmax projection, while approximation guarantees, generalization bounds, and compression-error theory remain open directions.

Impact Statement

The work presents functional attention as a formulation intended to improve operator-learning robustness, efficiency, and generalization across resolutions and geometries. It targets scientific and engineering applications involving PDEs and geometric data.

  • Functional attention improves robustness, efficiency, and generalization across resolutions and geometries in operator learning.
  • The approach may benefit surrogate models for PDEs and geometric data in physical simulation, design optimization, and data-efficient modeling.

A.1. Proof of Proposition 4.3

The appendix analyzes Functional Attention’s basis behavior, integral-operator interpretation, stability, and relationships to existing attention formulations. These results connect its learned functional representation to established operators and limiting cases.

  • As temperature approaches zero, the softmax basis converges to a hard partition that recovers classical P0 piecewise-constant elements.
  • Functional Attention approximates a learnable integral operator on the domain with a learnable kernel.
  • The functional-attention computation corresponds to a Monte-Carlo discretization of the integral operator.
  • The method’s differential is bounded through term-by-term norm estimates, with constants depending polynomially on model norms and softmax-Lipschitz factors.
  • With any orthonormal basis spanning the full space, Functional Attention reduces to Intention as a special case.

B. Complexity Analysis

FUNCATTN reduces attention’s dependence on sequence length by projecting through adaptive bases and solving a small latent transport system. Its complexity is linear in sequence length, and benchmarks show favorable runtime and memory at large resolutions.

  • Adaptive basis computation costs O(ndk) for the linear transformation and O(nk) for softmax, and must be recomputed when the input changes.
  • Latent projections have complexity O(ndk), while the transport solve operates on small matrices because k and d are much smaller than n.
  • The Woodbury reformulation permits inversion of either a d × d or k × k matrix with numerically identical results.
  • O(ndk + dk min(k, d) + min(k, d)^3) is the stated FUNCATTN complexity, linear in n and d and cubic in the typically small k.
  • At large sequence lengths, FUNCATTN has the smallest wall-clock time and memory footprint among tested attention variants.

C.2. PDE Benchmarks

The PDE evaluation spans eight benchmarks covering diverse equations, geometries, discretizations, and physical scenarios. FUNCATTN is assessed with relative L2 error and implementation choices including basis count and projection stability.

  • Benchmark suite: Eight PDE benchmarks cover fluid dynamics, elasticity, plasticity, porous-media flow, airfoil flow, pipe flow, and irregular geometric settings.The suite includes Burgers, Navier–Stokes, Elasticity, Plasticity, Airfoil, Pipe, Darcy, and Darcy Flow with Notch.
  • Benchmark suite: The tasks vary in inputs and outputs, including initial conditions, geometry, permeability, boundary conditions, punch geometry, and mesh point locations.Outputs include solution fields, displacement evolution, pressure fields, velocity fields, and Mach-number fields.
  • Evaluation: Relative L2 error is the primary evaluation metric, with AirfRANS additionally reporting drag and lift coefficient errors and Spearman correlation.Spearman correlation measures preservation of design rankings across test samples.
  • Training configuration: The common architecture uses eight transformer layers and eight attention heads, with larger hidden dimensions for Navier–Stokes and AirfRANS.Most benchmarks use hidden dimension 128, while Navier–Stokes and AirfRANS use 256; basis counts are reduced for those higher-complexity tasks.
  • Ablation: Moderate mode counts generally balance expressiveness and generalization, but the optimal number of bases varies across tasks.Elasticity and Plasticity favor 256 bases, while Darcy benefits from higher counts, likely reflecting differences in solution smoothness.
  • Projection choice: Using the transpose projection instead of the Moore–Penrose pseudo-inverse avoids exploding gradients and provides stable training, lower computational cost, and better accuracy.A Tikhonov-stabilized pseudo-inverse resolves instability but increases the inverted matrix’s condition number by more than an order of magnitude.

D.3. Sensitivity to Tikhonov Parameter λ

The sensitivity study treats the Tikhonov parameter λ as a numerical stabilizer and examines how its initialization affects conditioning during Elasticity training. FUNCATTN remains accurate across the tested positive-regularization settings.

  • Experimental setup: λ is learnable through α and is varied by changing its initialization to study regularization strength during training.The study tracks the average condition number across the eight FUNCATTN layers on Elasticity.
  • Sensitivity results: Weaker regularization raises the final condition number from ∼8 to ∼100, while Elasticity test error varies by less than 0.02 across three settings.The reported comparison appears in Table 13 for different α initializations.
  • Interpretation: The authors conclude that λ must remain strictly positive for stability, whereas its exact value is not a sensitive hyperparameter within the tested range.This conclusion is linked to instability as λ approaches zero and the 1/λ2 scaling stated in Proposition 4.5.
  • Basis visualization: FUNCATTN learns smooth, localized basis functions, while orthogonality constraints produce globally supported bases resembling Fourier modes.The comparison is presented as a visualization of learned bases for different models.

E.2. PDE Visualization

The visualization section compares FUNCATTN and Transolver through spatial error maps and predicted fields across PDE benchmarks. These views complement scalar relative L2 metrics by showing where models succeed or struggle.

  • Error maps: Qualitative comparisons across six benchmarks use absolute error maps to reveal spatial error distributions alongside scalar relative L2 metrics.The visualizations identify regions where models struggle, including boundaries and areas with sharp gradients.
  • Comparison layout: The figures compare ground truth, Transolver predictions, FUNCATTN predictions, and corresponding error maps.The comparison is organized as visual field outputs and spatial error distributions.
  • PDE examples: Prediction visualizations include Darcy flow solutions and Elasticity stress fields on irregular meshes.These examples emphasize field prediction across different physical quantities and geometries.
  • PDE examples: Additional visualizations show Airfoil velocity fields, Navier–Stokes vorticity after rollout, Plasticity displacement magnitude, and Pipe flow velocity fields.The examples span steady fields, temporal rollouts, displacement magnitudes, and irregular-mesh flow predictions.
Loading 2605.31559v1…