Proposal: HEALPix field-space attention as a Samudra backbone

[hammer]

Proposal: a HEALPix field-space attention backbone for Samudra

(Architecture proposal — relates to the "Next architecture ideas" tracker #422, and to the axial-attention idea in #600.)

TL;DR

Add an optional field-space attention backbone — a structure-preserving Earth-system transformer that operates directly on ocean fields on an equal-area HEALPix sphere — as an alternative to the ConvNeXt-U-Net / Perceiver backbones. It's attractive because it (1) keeps every layer's state a real ocean field (no latent bottleneck), (2) does dense attention only over coarse HEALPix cells while carrying multiscale detail in the channel dimension (so it sidesteps the ~65k-token cost of global attention on the 1°×1° grid), (3) is parameter-efficient, and (4) removes the lat–lon pole distortion via equal-area cells. There is a permissively-licensed PyTorch reference (CC0) to port from, and it slots into Samudra's existing config-selected, dependency-injected backbone pattern. The one genuinely new piece is making the method's multiscale decomposition land/bathymetry-aware (relates to #303).

The methods

• Field-Space Attention for Structure-Preserving Earth System Transformers — Witte, Meuer, Plésiat, Kadow (DKRZ), arXiv 2512.20350.
• Field-Space Autoencoder for Scalable Climate Emulators — Meuer, Witte, Plésiat, Ludwig, Kadow, arXiv 2601.15102 (the generative/compression follow-up).
• Reference implementation: FREVA-CLINT/FieldSpaceNN — PyTorch (Lightning), CC0-1.0, covering both papers (field-space attention modules, autoencoder, diffusion, HEALPix support, CNN/CNN-VAE baselines).

How it works, briefly. Fields live on the HEALPix sphere throughout; each block is a residual, structure-preserving update x_{n+1} = x_n + Δ_θ(x_n), so intermediate states stay interpretable fields. A fixed, non-learned multiscale decomposition (coarse mean + detail residuals, a Laplacian-pyramid analog on the HEALPix quadtree) replaces learned patchify; the multiscale detail is stacked into the channel dimension and the attention tokens are the coarse HEALPix cells. The attention itself is then standard scaled-dot-product MHA (with pre-LN + AdaLN/FiLM + MLP) — with a single zoom level it reduces exactly to a plain ViT layer. Positional encoding uses spherical-harmonic init band-limited to the grid Nyquist degree. A "scale-conservation" step forces detail residuals to zero mean within each parent cell (a spatial-scale conservation prior). In the source paper a ~15.5M-param field-space transformer beat a 214M multiscale-ViT on global temperature super-resolution, with faster/more stable convergence.

Why this could help Samudra specifically

Motivated by the weaknesses documented in the Samudra papers (2412.03795, 2606.02610): the ~20–50% underestimation of forcing/climate-change trends, and the long-rollout variance collapse / imprinting artifacts that Samudra 2 addressed. Field-space attention's design speaks to these:

• Token count is solved by construction. The 1°×1° grid is ~65k cells — too many for dense global attention. Samudra-Multi handles this by compressing to a few hundred Perceiver latents, which discards spatial structure. Field-space attention instead attends over a small number of coarse cells (e.g. a coarse HEALPix level → ~hundreds–low-thousands of tokens) while keeping fine detail in channels and every layer a real field — a strictly better structural match for a field-to-field emulator.
• Equal-area, no pole distortion. Samudra's loss is currently a uniform (non-area-weighted) mean over lat–lon, which over-weights the poles; HEALPix equal-area cells make uniform weighting correct and drop the zero-padded-pole artifact.
• Parameter efficiency — appealing for a model trained on ~4×A100.
• The scale-conservation prior is at least thematically aligned with the drift/variance issues (see honest caveat below).

Why Samudra is a good host (fits the existing seams)

From reading the current main (Samudra-2-era package):

• Backbones are config-selected and dependency-injected. The model class is a Pydantic discriminated union (AnyModelConfig in src/samudra/config.py) and backbones are built via build() factories and passed into the model (Samudra(unet=...)). Adding a field-space backbone = a new *Config(BaseModelConfig) + build() returning an nn.Module, without touching the trainer or data path.
• Attention + sphere-native coordinates already exist. src/samudra/models/samudra_multi.py already uses perceiver-pytorch (+ optional flash-attn), and its Perceiver-IO decoder queries each output pixel by its continuous 3D unit-sphere coordinate (via models/modules/augment_input.py / Aurora positional encodings) explicitly to generalize across resolutions. So resolution-agnostic, coordinate-based decoding is already in the repo and largely reusable.
• The regrid tooling is already a dependency. xesmf is in pyproject.toml, so the lat–lon → HEALPix (or ideally OM4-native → HEALPix) conservative remap is a preprocessing step in data/ocean_preprocessing/, not new infra.
• PyTorch + Apache-2.0, and the reference code is PyTorch + CC0 — so the field-space attention modules can be lifted in directly.

Proposed implementation (phased)

1. Preprocessing. Add a HEALPix regrid of the OM4 fields and the per-(variable, depth) "wet" mask via xesmf; choose nside. Emit tokens (B, N_pix, C) and mask (vars, N_pix).
2. Mask/context plumbing. Reshape the Masks / GridContext from (…, lat, lon) to (…, N_pix) (utils/data.py, utils/ctx.py, utils/loss.py). Loss stays a masked mean — now equal-area-justified.
3. Backbone module. New AnyModelConfig member + nn.Module porting FieldSpaceNN's field-space attention block: fixed HEALPix multiscale decomposition, child→channel reshape, scale conservation, MHA/MLP/AdaLN over coarse-cell tokens, spherical-harmonic positional init. The attention core is standard transformer code; the novel part is the HEALPix index arithmetic (already implemented in the reference repo).
4. Autoregressive wiring. Map the residual field update x_{n+1}=x_n+Δ onto forward_once/rollout. Add an input projection to tame the channel-width interaction (Samudra's ~150 channels × the multiscale channel-expansion), and choose zoom levels / attention level deliberately.
5. A/B evaluation. Same OM4 data; swap SamudraConfig/SamudraMultiConfig for a SamudraFieldSpaceConfig; compare rollout skill, long-horizon stability, and variance retention.

The one genuinely new piece: land-mask-aware multiscale (relates to #303)

The field-space papers assume full spherical coverage (atmosphere/global temperature). The ocean has coastlines and depth-varying bathymetry — Samudra's wet mask is effectively 19 nested ocean domains. So the multiscale decomposition's parent-cell averaging must become a masked mean over wet children only, scale-conservation's "zero-mean within a parent" must be "zero-mean over the wet children," partially-wet coarse cells need a policy, and attention should mask land tokens. This is the part that needs real design + validation and doesn't come for free from the reference code. It also connects to the land-boundary work proposed in #303 and the grid-handling in #801.

What's reused vs. built

• Reused: PyTorch trainer, rollout curriculum, physical corrector.py, EMA, Pydantic/YAML config flow, xesmf regridding, the Aurora/3D-coordinate positional machinery, and the FieldSpaceNN attention modules (CC0).
• Built: the HEALPix regrid + masked multiscale/scale-conservation for bathymetry, the (…, N_pix) token/mask plumbing, and the backbone config/module.

Scope options

• Now — field-space attention as a backbone (this proposal): a bounded project; most effort is the masked-HEALPix multiscale and validation.
• Later — the field-space autoencoder + diffusion (2601.15102): would give Samudra a probabilistic/generative mode (ocean ensembles, uncertainty, compression, zero-shot super-resolution by masking fine residual tokens) — directly targeting variance collapse and ensemble use cases. Much larger lift; noting it as a direction, not part of this issue.

Honest caveats

• Field-space "conservation" is a spatial-scale conservation (zero-mean residuals across scales), not physical heat/salt/mass conservation — it's a plausible help for drift/variance, not a guaranteed fix; corrector.py remains the real physical enforcer.
• The masked multiscale over bathymetry is un-tried (the papers don't cover land/ocean boundaries) — the main technical risk.
• Going lat–lon → HEALPix is a second lossy remap after OM4→lat–lon; regridding OM4-native → HEALPix directly would avoid double interpolation (more preprocessing, cleaner).
• HEALPix removes the exact zonal periodicity the conv backbone relies on (circular longitude padding) — a different inductive bias to validate.

Open questions for maintainers

1. Is a HEALPix representation of interest for Samudra, or is staying on the native/curvilinear grid preferred (cf. Untangle i/j/x/y/lat/lon given LLC datasets too #801 grid discussion)?
2. Preferred nside / target token budget, and appetite for the direct OM4→HEALPix regrid vs. the existing lat–lon intermediate?
3. Would you rather see this as a backbone A/B first, or is the generative field-space autoencoder direction the more interesting one for your ensemble/uncertainty goals?

Happy to prototype the SamudraFieldSpaceConfig backbone module and the masked-HEALPix-multiscale function against OM4 if this is a direction you'd welcome.

References

Field-Space Attention — arXiv 2512.20350 · Field-Space Autoencoder — arXiv 2601.15102 · reference code FREVA-CLINT/FieldSpaceNN (CC0) · Samudra — arXiv 2412.03795 · Samudra 2 — arXiv 2606.02610.