Function-Space Transformer with Adaptive Anchors

Guorui Sang; Pedram Rooshenas

cs.LGSep 29, 2026

Function-Space Transformer with Adaptive Anchors

Authors: Guorui Sang, Pedram Rooshenas

Organizations: University of Illinois Chicago

Abstract

Many forms of data, including physical fields, geometric shapes, and visual signals, are naturally described by functions over continuous domains but are observed through discrete samples. Representing these functions on fixed uniform grids imposes a trade-off between resolving localized variation and increasing computation across the domain. Neural operators address this mismatch by learning mappings between functions, while latent-attention architectures provide flexible processing of sampled observations. We introduce the Function-Space Transformer (FST), a framework for learning from functions through a spatially adaptive continuous latent representation. FST stores features at anchors whose locations are predicted from the input observations and recursively refines these anchor features through function-space interactions. This allows the representation to adapt its spatial organization to each input rather than inherit that of the observation grid, while supporting both spatially resolved and finite-dimensional outputs. On PDE solution prediction using PDEBench Burgers and Darcy flow, FST substantially outperforms the Perceiver IO baseline, whose latent representation lacks explicit spatial organization, and is highly competitive with the Fourier Neural Operator. On ImageNet-1K, FST achieves higher classification accuracy than the Vision Transformer baseline, with fewer parameters across these comparisons. Ablations further support the benefits of function-space updates and recursive refinement. Together, these results highlight the potential of adaptive continuous representations for both scientific prediction and visual recognition.

Figures & tables

Model	Latent representation	Global interaction
Perceiver IO / UPT	Unstructured latent tokens	Self-attention among latent tokens
ENF	Latent features optimized for each input with Gaussian windows	Message passing among latents (ENF-PDE)
GPO	Predicted Gaussian geometry	Attention among modal tokens
AB-UPT	Spatial anchor tokens; separate query tokens	Self-attention among anchors; queries only read anchors
FST (ours)	Anchors predicted from the input, normalized RBFs	Attention among sampled field evaluations, written back

Table 1: Where each architecture stores its latent representation and where it computes global interaction.

Figure 1: Function-Space Transformer (FST). (a) A shared layer refines anchor features for

R

steps at fixed coordinates

\mathbf{X}_{A}

, predicted once from the input. The decoder queries the field at

\mathcal{Q}

or pools anchor features. (b) Each step attends to observations, evaluates RBF blends through a pointwise MLP at

\mathbf{X}_{E}^{(r)}

, and applies self-attention. Cross-attention updates anchors using queries from the same RBF blend and MLP at

\mathbf{X}_{A}

Figure 2: Evaluate, mix, and write-back. (a) RBF weights blend anchor features; a pointwise MLP, omitted for clarity, maps the blend and coordinate embedding to an evaluated feature. (b) Evaluation points may include anchor coordinates. (c) Self-attention mixes evaluated features. (d) Cross-attention incorporates information into anchor features. Its queries use the same RBF blend and MLP at anchor coordinates. Anchor coordinates stay fixed; one receiving anchor is highlighted.

Model	Supervision per step	Parameters	Stop step	Relative $L_{2}$
FST	2,500 points	2.901M	500,000	0.018180
FNO	full grid	2.978M	400,000	0.018528
Perceiver IO	2,500 points	2.995M	900,000	0.073293

Table 2: Burgers results across all 12 viscosities. Stop step: end of training. Selected steps (lowest validation error): FST 450k, FNO 400k, Perceiver IO 895k.

Figure 3: Per-viscosity test error on PDEBench Burgers (all-viscosity models). (a) Relative

L_{2}

error. (b)

100\,(e_{\mathrm{FST}}-e_{\mathrm{FNO}})/e_{\mathrm{FNO}}

; negative values favor FST. All 12 viscosities are shown.

Figure 4: Burgers profiles at

\nu=0.001

(top) and

\nu=0.02

(bottom). Colors indicate time; each row shares axis limits. FST captures fronts with fewer spurious oscillations (top) but rounds initial-state minima (bottom).

Model	Parameters	Supervision per step	$32\times 32$	$64\times 64$	$128\times 128$	$201\times 1024$
FST	2.831M	2,500 points	0.04118	0.04159	0.04171	0.04185
FNO	2.978M	full grid	0.02854	0.02950	0.02963	0.02974
Perceiver IO	2.937M	2,500 points	0.25634	0.25465	0.25396	0.25381
Transolver	2.956M	2,500 points	0.17709	0.07151	0.06569	0.06154
DeepONet	2.998M	2,500 points	0.24712	0.25204	0.25467	0.25581

Table 3: Test relative

L_{2}

error on Burgers (

\nu=0.01

) at four output grids. See Appendix B for detailed inputs, outputs, and evaluation.

Coordinates	Length scales	$32^{2}$ anchors	$16^{2}$ anchors
Uniform	Fixed	0.04604	0.04571
Adaptive	Fixed	0.04069	0.05784
Uniform	Adaptive	0.05482	0.05094
Adaptive	Adaptive	0.04185	0.04480

Table 4: Ablation on anchor coordinates and length scales.

Table 9

Table 10

Figure 5: Full-grid relative

L_{2}

error using different numbers of nearby anchors.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Converting predicted gaps to anchor coordinates. Each row has its own predicted spacing.

Setting	Burgers	Darcy	ImageNet-1K
Anchors $M$	$32\times 32=1{,}024$	$32\times 32=1{,}024$	$32\times 32=1{,}024$
Anchor feature dimension	176	176	200
Main attention heads	8	8	8
Observation-attention bias	Zero	Zero	Eq. 11
Refinement steps $R$	4	4	4
Evaluation points $S$	512	512	512

Appendix

Table 9: FST architecture settings for the reported experiments.

Component	Burgers	Darcy	ImageNet-1K
Coordinate predictor	Fourier coordinate embeddings	Fourier coordinate embeddings	Fourier embeddings of patch-center coordinates
Predictor observation inputs	Coordinates, values, viscosity	Coordinates, coefficient values	Patch centers, patch features

Appendix

Table 10: Coordinate predictor encodings and observation inputs.

Setting	All viscosities	$\nu=0.01$
Training samples	24,000	2,000
Validation samples	1,200	100
Test samples	1,000	84
Original grid (time $\times$ space)	$201\times 1024$
IC values	1,024
BC values	$2\times 200=400$

Appendix

Table 11: Burgers datasets and shared observation grid.

Figure 7: Burgers decoding at one query coordinate.

Model	Input	Supervision per training step
FST	ICs and BCs + $\nu$	2,500 sampled points
Perceiver IO	ICs and BCs + $\nu$	2,500 sampled points
FNO	Values + mask + $\nu$ on the $201\times 1024$ grid	Full grid (205,824 points)

Appendix

Table 12: Model inputs and output supervision for the Burgers comparison.

Model	Width	Layers	Heads	Latents/slices	Modes
FNO	48	4	-	-	$10\times 30$
Perceiver IO	240	6	8	256	-
Transolver	328	5	8	32	-

Appendix

Table 13: Burgers baseline architecture settings. Fourier modes follow the time-space axis order.

Experiment	Anchors	Coordinates	Length scales	$R$
All viscosities	$32^{2}$	Adaptive	Adaptive	4
Single viscosity	$32^{2}$	Adaptive	Adaptive	4
Default ablation	$32^{2}$	Adaptive	0.03	4

Appendix

Table 14: FST configurations for the Burgers comparisons and the default ablation. The default settings are varied in Tables 4 - 6 .

Setting	Value
Optimizer	Adam
Global batch size	48
Forward / backward precision	bfloat16 (BF16)
Loss accumulation precision	32-bit floating point (FP32)
Minimum training steps	250,000
Maximum training steps	1,000,000

Appendix

Table 15: Shared training settings for the all-viscosity Burgers comparison.

Model	Selected step	Stop step
FST	92,000	98,000
FNO	72,000	72,000
Perceiver IO	34,000	50,000
Transolver	80,000	82,000
DeepONet	32,000	48,000

Appendix

Table 16: Checkpoints for the single-viscosity models in Table 3 .

Coordinates	Scales	$32^{2}$	$16^{2}$
Uniform	Fixed	82k	116k
Adaptive	Fixed	86k	80k
Uniform	Adaptive	130k	112k
Adaptive	Adaptive	92k	78k

Appendix

Table 17: Selected training steps for the Burgers ablations; k denotes 1,000 steps. Left: coordinate and length-scale ablations at two anchor counts. Right: update and refinement ablations with

32^{2}

anchors.

Model	GPUs	Time (h)	Memory (GiB)
FST	1	30.5	8.80
FNO	3	32.9	15.21
Perceiver IO	3	18.6	2.32

Appendix

Table 18: All-viscosity training on NVIDIA A100 80GB GPUs. Time sums the recorded training segments, including validation. Memory is the peak allocated memory per GPU.

$\nu$	FST	FNO	Perceiver IO
0.001	0.044195	0.048623	0.147162
0.002	0.033692	0.036266	0.120887
0.004	0.027352	0.031828	0.112725
0.01	0.021308	0.024882	0.088623
0.02	0.013957	0.015056	0.063128
0.04	0.013382	0.013013	0.061419

Appendix

Table 19: Per-viscosity relative

L_{2}

error on PDEBench Burgers. Best in bold; second best underlined.

Figure 8: Additional Burgers solution profiles. From top to bottom, rows show selected trajectories at $\nu=0.001$ , $0.002$ , $0.004$ , $0.01$ , and $0.02$ . Columns show the reference solution and predictions from FST, FNO, and Perceiver IO. Colors indicate $t\in\{0,0.1,0.3,0.6,1,2\}$ ; axis limits are shared across columns within each row. The first and last rows are the examples in Figure 4 . Different rows have different initial conditions.

Figure 9: Burgers solution profiles for different initial conditions at fixed viscosity $\nu=0.002$ . From top to bottom, the initial profiles are Gaussian-like, single wave, two waves, and localized oscillations. Columns show the reference solution and predictions from FST, FNO, and Perceiver IO. Colors indicate $t\in\{0,0.1,0.3,0.6,1,2\}$ ; axis limits are shared across columns within each row. The examples contrast broad spatial features, repeated wave structures, and spatially concentrated oscillations.

Split	Samples	Grid	Precision
Training	2,000	$128\times 128$	BF16
Validation	100	$100\times 100$	FP32
Test	1,000	$128\times 128$	FP32

Appendix

Table 20: Darcy data and training settings.

Figure 10: Darcy anchor coordinates and decoding RBF length scales for four test samples. (a) Anchors over reference solutions. (b) Soft ellipses with horizontal and vertical scales proportional to $\ell_{x}$ and $\ell_{y}$ , displayed at $0.35\times$ scale; color indicates $\sqrt{\ell_{x}\ell_{y}}$ .

Setting	Value
Training / validation images	1,281,167 / 50,000
Training crop	Random resized, $224\times 224$
Training augmentations	Horizontal flips, 3-Augment ( Touvron et al., 2022 )
Color jitter	0.3
Random erasing probability	0.25
Validation resize (short side)	256

Appendix

Table 21: ImageNet-1K and data preprocessing.

Stage	Operation	Setting
1	Convolution	64 channels
2	Residual blocks	2 blocks, 2 convolutions each
3	Pointwise projection	Width 200
4	Average pooling	Within each patch
5	Add positional embedding	Fourier embeddings of patch bounds + MLP
6	Residual feed-forward network	-

Appendix

Table 22: FST patch encoder, in execution order. Spatial convolutions use

3\times 3

kernels and the encoder uses GELU activations.

Figure 11: ImageNet patch observations. A CNN encodes each $16\times 16$ patch. Spatial pooling, a patch-bound embedding, and a residual FFN produce the feature $\mathbf{h}_{i}$ , paired with the patch center $\mathbf{c}_{i}=((x_{i}^{0}+x_{i}^{1})/2,(y_{i}^{0}+y_{i}^{1})/2)$ .

Setting	Value
Patch projection	Linear, $16\times 16$ patches
Feature width	176
Transformer layers	10
Attention heads	4
Feed-forward hidden width	423
Trainable parameters	3,095,614

Appendix

Table 23: ViT baseline configuration.

Setting	Value
Optimizer	AdamW
Weight decay	0.05
Effective batch size	256
Computation precision	BF16
Gradient norm clipping	1.0
Label smoothing	0.1

Appendix

Table 24: Shared ImageNet-1K optimization settings.

Stage	Start	End	Learning rate
Warmup	0	25,023	Peak $10^{-3}$
Cosine decay	25,023	500,456	$10^{-3}\to 5\times 10^{-5}$
Cosine tail	500,456	550,456	$5\times 10^{-5}\to 0$
Constant learning rate	550,456	600,456	$2.5\times 10^{-6}$

Appendix

Table 25: Learning-rate schedule for FST and ViT.

Figure 12: Anchor coordinates and RBF length scales on ImageNet-1K validation images. (a, c) All 1,024 anchors overlaid on the images. (b, d) RBF length scales $\ell^{\mathrm{ref}}$ , shown by color and soft circles displayed at $0.35\times$ scale. Horizontal coordinates adapt to each input; vertical rows are fixed.

Explore similar work

Functional Attention: From Pairwise Affinities to Functional Correspondences
May 29, 2026Jiefang Xiao, Maolin Gao, Simon Weber +2Transformer Attention Dynamic Attention
Technical University of Munich, Germany · Munich Center for Machine Learning (MCML), Germany · PIXL, Department of Computer Science, University of Oxford, United Kingdom +1
Function graph transformers universally approximate operators between function spaces
May 18, 2026Takashi Furuya, David Mis, Ivan Dokmanić +2Probabilistic Operator Learning Multiplex Graph Transformers
Doshisha University, RIKEN AIP · Rice University · University of Basel +2
MoNo: Multiscale Optimal Transport Neural Operator for Solving PDEs on General Geometries
Aug 10, 2026Zijiang Yang, Xiaomeng Wu, Dongmei FuNeural Operators Neural Partial Differential Equation Solvers
School of Automation and Electrical Engineering, University of Science and Technology Beijing · 2Beijing Engineering Research Center of Industrial Spectrum Imaging

May 29, 2026, cs.LGJ/K move · Enter open · S save

Functional Attention: From Pairwise Affinities to Functional Correspondences

Jiefang Xiao, Maolin Gao, Simon Weber, Guandao Yang+1

Technical University of Munich, Germany · Munich Center for Machine Learning (MCML), Germany · PIXL, Department of Computer Science, University of Oxford, United Kingdom+1

Transformer Attention Dynamic Attention

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Functional Attention: From Pairwise Affinities to Functional Correspondences

Function graph transformers universally approximate operators between function spaces

MoNo: Multiscale Optimal Transport Neural Operator for Solving PDEs on General Geometries