cross-attn-fno-darcy

Cross-Attention Coefficient Head for OOD Darcy Flow

Novel contribution: A drop-in cross-attention input head for FNO/Geometric-DeepONet that treats the variable-coefficient field as keys and values rather than a concatenated channel. Benchmarked on shifted log-Gaussian Darcy flow.

Problem

Reservoir simulation surrogates are retrained per site because permeability statistics differ across reservoirs. The standard approach concatenates the coefficient field as an extra channel. We hypothesize that explicit cross-attention conditioning — where the spatial solution queries attend over coefficient keys/values — forces the model to build a robust conditioning representation that generalizes under distribution shift.

Architecture

FNOBaseline (Channel-Concat Standard)

a(x) ──[1x1 conv]──► width channels ──[FNO blocks]──► projection ──► u(x)

CrossAttnFNO (Our Contribution)

a(x) ──[kv_embed]──► KV
         │
         ├──► cross-attn ◄── Q = query_proj(coordinate_grid [-1,1]²)
         │
         └──► bypass = coeff_bypass(a) ──┐
                                          ▼
              attended + bypass ──[FNO blocks]──► projection ──► u(x)

Key design:

• Coordinate queries: Each spatial location queries what coefficient features matter locally
• GNOT-style feature normalization: Q/K normalized over embedding dim (not sequence dim)
• Bypass residual: Model can fall back to standard behavior
• O(N) complexity: Attention within sample, not across samples

OperatorSplittingFNO (Polymathic-Style Baseline)

MoE composition with learned soft gating from global coefficient statistics.

Dataset

• PDE: 2D Darcy -∇·(a(x)∇u) = 1 with zero Dirichlet BCs
• Coefficient: Log-Gaussian random field, isotropic covariance exp(-L²|k|²/2)
• Train: L=0.1, 1000 samples (full) / 200 samples (small)
• Test ID: L=0.1
• Test OOD smooth: L=0.2 (longer correlation)
• Test OOD rough: L=0.05 (shorter correlation)

Training Recipe

Quick Start

Small experiment (CPU, ~1-2 hours)

python train_small_v2.py
# Output: results_small.json, best_*.pt

Full experiment (GPU recommended, ~3-4 hours)

python train_full.py
# Output: results.json, best_*.pt

Analyze results

python analyze_results.py results_small.json
# Output: comparison.png, degradation.png

Files

Literature

• Li et al. FNO (NeurIPS 2021): Baseline architecture
• GNOT (2302.14376): Heterogeneous attention for multi-input operators
• Continuum Attention (2406.06486): Theoretical justification
• SpectraKAN (2602.05187): Closest architectural cousin — single-query cross-attention in spectral operators
• Polymathic Operator Splitting (2602.00884): Test-time correction baseline
• ANCHOR (2512.19643): EMA residual triggering
• Fundamental Physics (2602.15184, ICLR 2026): Multi-task decomposition

Falsification Criteria

1. Cross-attn degradation ≤ 1.5× on shifted permeability
2. Channel-concat baseline degradation ≥ 3× on shifted permeability
3. Operator-splitting at matched compute: wins or loses by measurable margin

Status

• ✅ Architecture implemented (3 models)
• ✅ Data generation pipeline (log-Gaussian + Darcy solver)
• ✅ Solver verified (no NaN targets)
• 🔄 Training in progress (Job: https://huggingface.co/jobs/jmtsh21/69fecb03317220dbbd1a7098)

License

MIT