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

Architecture

FNOBaseline (Channel-Concat Standard)

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)

• 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)

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)

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

Analyze results

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