Geometry-Aware Operator Families for Structured Representation Learning

Zuyuan Zhang; Fei Xu Yu; Tian Lan

cs.LGSep 27, 2026

Geometry-Aware Operator Families for Structured Representation Learning

Authors: Zuyuan Zhang, Fei Xu Yu, Tian Lan

Organizations: The George Washington University

Abstract

The geometry of latent representations governs which components should interact and how information should propagate, making geometry-aware operator design a fundamental ingredient of structured deep representation learning. However, existing neural architectures typically rely on generic operator templates or geometry-specific constructions, creating a need for a unified framework that can derive admissible operators directly from fixed structural information while remaining adaptive to changing contexts. We introduce Geometry-Induced Operator Families (GIOF), a general framework that converts fixed geometry into a structured family of propagation operators and dynamically selects an appropriate member of this family according to the current context. GIOF first transforms geometry-derived interaction channels into reusable generator bases, then combines them through a context-dependent selector and adaptive propagation scale, and finally realizes the selected operator through stable continuous-time propagation and a bottleneck residual layer. We establish theoretical guarantees covering parameter compression, identifiability, stability, locality, compositional structure, and oversmoothing behavior, while controlled experiments validate these mechanisms and experiments on PEMS-BAY and METR-LA achieve the lowest mean MAE across all reported regional-outage settings, improving over the strongest retained baseline by 2.4%--8.8% at 30% missing sensors.

Figures & tables

Panel of Figure 1 — Figure 1: Controlled numerical checks of the theoretical mechanisms. (a) Structured probe recovery improves rapidly with the number of vector probes, while less structured estimators require substantially more observations. (b) Under matched geometry, the structured estimator approaches the exact projection floor, whereas increasing mismatch exposes the corresponding approximation cost. (c) Positive-rate paths exhibit finite-scale influence controlled by the corresponding Poisson-tail envelopes, with disconnected leakage remaining numerically zero. (d) Propagation reduces risk at intermediate scales for spectrally aligned signals, while excessive smoothing eventually increases risk.

Panel of Figure 2 — Figure 2: Robustness to regional outages with 20% training days. Test MAE is reported under increasing fractions of missing sensors using fixed checkpoints. The 60% missingness setting lies outside the mask rates used during training.

Method	20%	100%	20%	100%	Params (K)	Train (min)	GPU (MiB)	Infer (ms)
	PEMS-BAY MAE		METR-LA MAE		Cost on PEMS-BAY
GCN	$3.232\pm 0.051$	$3.208\pm 0.020$	$4.972\pm 0.056$	$5.229\pm 0.081$	18.98	1.04	29.95	1.02
GATv2	$3.853\pm 0.801$	$3.202\pm 0.060$	$4.964\pm 0.024$	$5.044\pm 0.018$	19.52	1.54	51.99	1.86
APPNP	$3.291\pm 0.036$	$3.280\pm 0.043$	$4.966\pm 0.022$	$5.091\pm 0.022$	18.43	1.56	47.58	2.20
MoNet	$3.223\pm 0.027$	$3.197\pm 0.027$	$4.984\pm 0.055$	$5.234\pm 0.122$	23.42	3.46	107.63	5.37
ECC	$3.198\pm 0.019$	$3.147\pm 0.035$	$4.897\pm 0.024$	$5.133\pm 0.080$	36.64	1.42	49.28	1.96

Table 1: Regional-outage reconstruction at 30% missing sensors and measured computational cost. MAE is in the original speed units; 20% and 100% denote the fraction of available training days. Values first average ten shared mask realizations within each model seed and then report mean

\pm

sample standard deviation. Cost columns use PEMS-BAY with 100% training days and average the three final runs. All methods use the same final update cap; training time includes validation.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

\eta(c)=\operatorname{softmax}(g_{\eta}(c))

Appendix

Algorithm 1 Geometry-Aware Balanced Transport Layer

A(c)=\sum_{r=1}^{R}\theta_{r}(c)\Phi_{r}+\operatorname{Diag}(\delta(c))

Appendix

Algorithm 2 Geometry-Aware Residual Adapter

\alpha(c)=\operatorname{softmax}(f_{\theta}(c))

Appendix

Algorithm 3 Geometry-Aware Propagation Layer

$N$	Isotropic rank	Singular rank	Isotropic error	Singular error
2	72–72	40–40	$5.03\times 10^{-30}\pm 2.3\times 10^{-30}$	$0.386\pm 0.20$
4	72–72	40–40	$4.79\times 10^{-30}\pm 1.5\times 10^{-30}$	$0.386\pm 0.20$
8	72–72	40–40	$4.91\times 10^{-30}\pm 2.4\times 10^{-30}$	$0.386\pm 0.20$
16	72–72	40–40	$3.45\times 10^{-30}\pm 8.8\times 10^{-31}$	$0.386\pm 0.20$
32	72–72	40–40	$3.64\times 10^{-30}\pm 1.4\times 10^{-30}$	$0.386\pm 0.20$
64	72–72	40–40	$3.68\times 10^{-30}\pm 8.1\times 10^{-31}$	$0.386\pm 0.20$

Appendix

Table 2: Noiseless probe identifiability. Ranks show the range across seeds; errors are mean

\pm

sample standard deviation. Isotropic probes identify the full 72-dimensional dictionary, whereas the singular covariance leaves 32 diagonal directions unobserved and reduces the rank to 40.

Check	Maximum residual	Tolerance	Status
Mean-exit normalization	$1.11\times 10^{-16}$	$10^{-9}$	Pass
Pathwise Poisson bound	$0$	$10^{-9}$	Pass
Disconnected leakage	$0$	$10^{-9}$	Pass
Composition: directed invariants	$1.67\times 10^{-15}$	$10^{-9}$	Pass
Composition: symmetric invariants	$1.33\times 10^{-15}$	$10^{-9}$	Pass
Projection identity	$8.88\times 10^{-16}$	$10^{-9}$	Pass

Appendix

Table 3: Maximum numerical residual over all specified seeds and scales. Bound residuals report only positive violations; zero means that no violation was observed at the stated precision.

Figure 3: Additional controlled diagnostics for identifiability and structured projection. (a) Probe-Gram rank: isotropic probes identify the full structured dictionary, whereas singular covariance reduces the informative rank. (b) Probe-Gram conditioning: the smallest positive eigenvalue quantifies the conditioning of the identifiable probe system. (c) Noiseless recovery: isotropic probing recovers the structured operator to numerical precision, whereas singular probing leaves a persistent unidentifiable component. (d) Anisotropic projection: under covariance-weighted prediction risk, the optimal structured approximation need not coincide with the Frobenius projection.

Figure 4: Additional controlled diagnostics for propagation and sensitivity. (a) Noncommuting composition: adjacent depth-dependent generators have nonzero commutators, while the composed propagators retain the required structural invariants. (b) Oversmoothing: zero-mean signals contract consistently with the spectral-gap upper bound. (c) Scale sensitivity: the propagation derivative with respect to the diffusion scale remains below the corresponding theoretical bound. (d) Shape sensitivity: perturbations of the generator obey the predicted sensitivity bound, with vanishing sensitivity at

\tau=0

Panel of Figure 5 — Figure 5: Additional missingness controls using the 20%-training-day checkpoints. (a,b) Independent sensor outages. (c,d) Independent point missingness. Curves report test MAE and aggregate seeds after averaging shared mask realizations within each seed.

Panel of Figure 6 — Figure 6: Diagnostics for structured missingness. (a,b) Reconstruction MAE as the contiguous missing interval grows. (c,d) Regional-outage error stratified by graph distance from available sensors. These evaluations reuse trained checkpoints and require no additional model fitting.

Variant	Reg. MAE	Ind. MAE	Point MAE	Reg. RMSE	Params (K)
No-spatial	$3.696\pm 0.002$	$3.566\pm 0.003$	$1.280\pm 0.053$	$6.886\pm 0.018$	16.38
Static	$3.483\pm 0.003$	$2.891\pm 0.006$	$1.307\pm 0.037$	$6.444\pm 0.005$	18.44
Shuffled	$3.501\pm 0.013$	$2.944\pm 0.015$	$1.291\pm 0.029$	$6.458\pm 0.027$	23.14
One-step	$3.534\pm 0.010$	$2.885\pm 0.032$	$1.306\pm 0.042$	$6.562\pm 0.023$	23.14
GeoMLP	$3.390\pm 0.007$	$2.870\pm 0.005$	$1.292\pm 0.023$	$6.263\pm 0.018$	23.27
FreeEdge	$3.384\pm 0.138$	$2.774\pm 0.117$	$1.364\pm 0.055$	$6.770\pm 0.793$	61.24

Appendix

Table 4: Real-data ablations on PEMS-BAY with 20% training days at 30% missingness. Reg.: regional outages; Ind.: independent sensor outages; Point: independent point missingness. Values report mean

\pm

sample standard deviation over seeds.

Representation Learning

Representation learning develops data transformations that expose useful structure for prediction, transfer, generation, and control.

See all 813 papers

Do Neural Operators Forget Geometry? The Forgetting Hypothesis in Deep Operator Learning
May 7, 2026Yanming Xia, Angelica I. Aviles-RiveroNeural Network Generalization Neural Operators
Qiuzhen College, Tsinghua University, Beijing, China · Yau Mathematical Sciences Center, Tsinghua University, Beijing, China.
Diffusion Operator Geometry of Feedforward Representations
May 1, 2026Kanishka ReddyNeural Representation Geometry Feedforward Neural Networks
Department of Applied Mathematics University of Washington Seattle, WA
Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning
Jun 16, 2026Oriol Vendrell-Gallart, Nima Negarandeh, Ramin BostanabadUncertainty Quantification PDE Surrogate Modeling
Department of Mechanical and Aerospace Engineering University of California, Irvine Irvine, CA 92617

May 7, 2026J/K move · Enter open · S save

Do Neural Operators Forget Geometry? The Forgetting Hypothesis in Deep Operator Learning

Yanming Xia, Angelica I. Aviles-Rivero

Qiuzhen College, Tsinghua University, Beijing, China · Yau Mathematical Sciences Center, Tsinghua University, Beijing, China.

Neural Network Generalization Neural Operators

Open paper AskPDF arXiv

Neural operators perform well on structured domains, yet their behaviour on irregular geometries remains poorly understood. We show that this limitation is not merely an encoding issue, but a depth-wise failure mode inherent to deep operator architectures. We formalise the Geometric Forgetting Hypothesis: due to the Markovian structure of operator layers and their reliance on global mixing mechanisms, neural operators progressively lose access to domain geometry as depth increases. Using layer-wise geometric probing, we demonstrate that both spectral and attention-based operators systematically lose geometric fidelity. We show that this geometric forgetting degrades accuracy, stability, and generalisation. To counteract it, we introduce a lightweight geometry memory injection mechanism that restores geometric constraints at intermediate depths with minimal architectural overhead. This simple intervention consistently mitigates forgetting and exposes a geometric shortcut instability in transformer-based operators, revealing that geometric retention is a structural requirement rather than a design choice.

Abstract

Representation Learning

Explore similar work

Do Neural Operators Forget Geometry? The Forgetting Hypothesis in Deep Operator Learning

Diffusion Operator Geometry of Feedforward Representations

Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning