axial attention blocks by amogh-gulati · Pull Request #611 · m2lines/Samudra

amogh-gulati · 2026-03-02T16:43:45Z

This branch tracks the development of axial and full attention in unet blocks for samudra v1

training run with axial attention at bottleneck : https://wandb.ai/ocean_emulators/default/runs/ps6jz798

training run vanilla samudra v1: https://wandb.ai/ocean_emulators/default/runs/c7of0p5t

Full attention at bottleneck with bigger attention block : https://wandb.ai/ocean_emulators/default/runs/xku4fh3o

fix made:
Changed unetbackbone to capture encoder skip tensors at the end of each encoder stage instead of coreblock. Earlier skip paths were not using the output from the attention layer instead they were just being forwarded after the convnext layer.

This branch enables optional U-Net attention in two forms: axial attention and full attention. Axial attention applies self-attention separately along the two spatial axes, first treating each longitude column as a sequence over height and then each latitude row as a sequence over width. Two axis-wise passes are applied sequentially with residual connections.

Full attention instead flattens the whole H×W feature map into one sequence of spatial tokens and attends over all token pairs directly.

In both cases, the channel dimension is treated as the embedding dimension for each spatial token, and multi-head attention splits those channels before the outputs are projected back and added residually to the feature map.

The branch also leaves an space for models at the bottleneck for future transformer-like block families, so the bottleneck can later host attention, transformer, or MaxViT-style blocks without another config redesign.

alxmrs · 2026-03-09T17:41:00Z

Lmk when this is no longer a draft PR and I'll be happy to review it!

amogh-gulati · 2026-03-09T18:58:01Z

Lmk when this is no longer a draft PR and I'll be happy to review it!

should be good now!

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0f0549ffd2

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-03-09T18:59:59Z

+                decoder_attention_configs,
+                reversed(self.ch_width),
+                strict=True,


Build decoder attention with decoder output channel widths

The decoder attention blocks are instantiated with reversed(self.ch_width), which starts at the bottleneck width, but in UNetBackbone each decoder attention block is applied after a decoder core block whose output channels are already reduced. With the default widths [200, 250, 300, 400], the first decoder attention is built for 400 channels but receives a 300-channel tensor, causing an immediate runtime error (e.g., GroupNorm channel mismatch) whenever any non-final decoder attention entry is enabled.

Useful? React with 👍 / 👎.

This may be related to the block creation order that I pointed out above.

alxmrs

I think there is a mistake that Codex and I caught related to the order of creating blocks on the decode side of the unet. I also provide some design suggestions for this feature. In general, the Axial and full transformers implementation are excellent. Here are some notes to help improve the PR for the next iteration.

alxmrs · 2026-03-09T19:46:03Z

@@ -0,0 +1,42 @@
+# yaml-language-server: $schema=../schemas/SamudraConfig.json


Nit: We may have just renamed our Samudra version to v2.

I'm using the base v1 configs to run experiments with attention to keep the changes isolated, should I use v2 instead?

Yes, I think we should base these off of v2.

alxmrs · 2026-03-09T19:50:09Z

@@ -0,0 +1,3 @@
+# Samudra OM4 V1 — Attention
+
+Use [model.yaml](/scratch/ag11542/ocean/Ocean_Emulator/Ocean_Emulator/configs/samudra_om4_v1_attention/model.yaml) as the base example and [train.yaml](/scratch/ag11542/ocean/Ocean_Emulator/Ocean_Emulator/configs/samudra_om4_v1_attention/train.yaml) to launch training.


Would love to have a high level explanation for the experiments in the README. In fact, what you wrote in the PR comment should be fine to re-purpose for this.

alxmrs · 2026-03-09T19:51:32Z

+    References:
+        Axial Attention in Multidimensional Transformers (Ho et al., 2019)
+        https://arxiv.org/abs/1912.12180
+    """


This is a good docstring.

alxmrs · 2026-03-09T19:52:06Z

+        num_heads: int = 8,
+        qkv_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        axis: Literal["height", "width"] = "height",


Let's omit default arguments, if possible.

alxmrs · 2026-03-09T22:20:39Z

+    """Configuration for a single attention block in the U-Net."""
+
+    attention_type: AttentionType = Field(
+        default="axial",


🐑 Let's make "full" attention the default. WDYT?

alxmrs · 2026-03-09T22:22:29Z

+    transformer: TransformerBottleneckConfig | None = Field(
+        default=None,
+        description="Reserved transformer bottleneck settings for future implementation.",
+    )
+    maxvit: MaxViTBottleneckConfig | None = Field(
+        default=None,
+        description="Reserved MaxViT bottleneck settings for future implementation.",
+    )


Let's omit these for now. Happy to add notes or TODOs here that these will come later.

alxmrs · 2026-03-09T22:25:42Z

+                decoder_attention_configs,
+                reversed(self.ch_width),
+                strict=True,


This may be related to the block creation order that I pointed out above.

alxmrs · 2026-03-09T22:29:21Z

+BottleneckBlockType = Literal["attention", "transformer", "maxvit"]
+
+
+class AttentionBlockConfig(BaseConfig):


In an ideal world, this would be rolled into BlockConfig. Maybe, to accomplish this, we would update the BlockConfig to be a union type, so we could pass context specific arguments depending on if the block was a CNN or Attention.

I'm imagining something similar to how the ModelConfig is done today. This is a bit tricky, and I (or I bet, Jesse) would be happy to meet to try to work through the pattern with you on this. I admit, the implementation path for this is fuzzy in my head.

What if we model each stage as an ordered list of components plus an explicit skip-capture point, instead of fixing every stage to has the same internal ordering?
In the current CNN-style U-Net, the natural order is coreblock then scale change but for a ViT-style stage with patch embedding is tokenize first (scale change), then transform.

(maybe transform keeps the spatial size consistent and transition changes it)

encoder_stages:

skip_after_component: 1
components:

kind: transform
block_type: convnext

kind: transform
attention_type: axial

kind: transition
transition_type: avg_pool

skip_after_component: 0
components:

kind: transition
patch_size: 4
embed_dim: 256

kind: transform
block_type: transformer

but this makes the skip logic pretty messy

alternative, a shape-preserving StageBlock abstraction that can contain one or more stage-preserving operations, such as a core block, an attention block, or both, and then keep ViTStageBlock as a separate stage type for the cases where tokenization/downsampling is coupled with the stage computation to keep skips simplified, VIT is the only one where we change shape before the feature mixing instead of after. It also keeps the skip logic cleaner, since the backbone would still only need to reason about one declared stage output per level

Right now, the conv or ConvNeXt block is the main block for each U-Net level. The attention block is just an extra step added after that main block. For the experiments with each level as a conv-based stage with optional attention on top the current config seems to let us experiment with axial and full attention on different coarse resolutions and if we dont need attention it does not pollute the config.

I think Convnext + attention blocks will be a more natural fit for the coreblock moving forward (as they define a stage and the current attention block do not, they are just a optional bolt on) and for patchify + transformer based design a new model like fomo might be the most apt way to experiment

@alxmrs let me know what you think here, happy to discuss this more

Mind if we discuss this over video chat after our weekly meeting? I'd like to understand what you're proposing a bit better (and ask some questions).

… of the coreblock

alxmrs

Another review pass. Sorry I missed your responses from earlier this week.

alxmrs · 2026-04-10T21:37:54Z

@@ -0,0 +1,42 @@
+# yaml-language-server: $schema=../schemas/SamudraConfig.json


Yes, I think we should base these off of v2.

alxmrs · 2026-04-10T21:49:44Z

+        num_heads: int = 8,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,


Is it possible to have these values be set by our config system? I'd prefer not to use default kwargs.

alxmrs · 2026-04-10T21:52:03Z

+        return x
+
+
+class FullAttention(nn.Module):


In the future, I might slightly update this module to support both self and cross attention; since this PR doesn't need it, this is good for now.

alxmrs · 2026-04-10T21:55:40Z

+
+    def __init__(
+        self,
+        channels: int,


Happy to punt on in/out channels; this is now made more complicated with "prog_channels" and "boundary_channels" in a WIP feature.

alxmrs · 2026-04-10T22:00:09Z

+            if decoder_attention_blocks is not None:
+                attention_block = decoder_attention_blocks[i]
+                if attention_block is not None:
+                    layers.append(attention_block)


In my understanding, this will break the "U" shaped symmetry of the network. Mirroring means reversing the order. This is only coming up now because we only ever had one kind of block.

If there is a convention for not having a strict reflection in symmetry, then I'm ok with this; I do want to check first.

alxmrs · 2026-04-10T22:02:42Z

+BottleneckBlockType = Literal["attention", "transformer", "maxvit"]
+
+
+class AttentionBlockConfig(BaseConfig):


Mind if we discuss this over video chat after our weekly meeting? I'd like to understand what you're proposing a bit better (and ask some questions).

amogh-gulati marked this pull request as draft March 2, 2026 16:43

amogh-gulati force-pushed the axial_attention branch from 1f0973d to c287675 Compare March 9, 2026 15:40

amogh-gulati marked this pull request as ready for review March 9, 2026 18:56

amogh-gulati requested a review from alxmrs March 9, 2026 18:58

chatgpt-codex-connector Bot reviewed Mar 9, 2026

View reviewed changes

alxmrs requested changes Mar 9, 2026

View reviewed changes

Amogh Gulati and others added 15 commits April 6, 2026 09:59

axial attention blocks

af103de

attention blocks and bottleneck blocks wiring

0413863

remove redundancy from bottleneck block

a4b70d3

Stop tracking experimental attention configs

6cbcebf

configs

c51e530

formatting

870f5f1

fix attention decoder logic

6cd355a

update readme

f6842e2

remove default args

b19d86c

Added jax typing

f6e36cb

remove unimplemented blocks

7f3dc6c

add an assert for attention heads

683577a

change default attention type to full

e662477

reshapes using einops

2e1ba88

make bottleneck block replace the default middle block

8c23fc2

amogh-gulati force-pushed the axial_attention branch from 88ab902 to 8c23fc2 Compare April 6, 2026 14:00

changed unet to capture the skip tensor the end of each stage instead…

87c1ef4

… of the coreblock

alxmrs reviewed Apr 10, 2026

View reviewed changes

alxmrs added this to Samudra Planning Aug 4, 2026

github-project-automation Bot moved this to Backlog in Samudra Planning Aug 4, 2026

		@@ -0,0 +1,42 @@
		# yaml-language-server: $schema=../schemas/SamudraConfig.json

		@@ -0,0 +1,3 @@
		# Samudra OM4 V1 — Attention

		Use [model.yaml](/scratch/ag11542/ocean/Ocean_Emulator/Ocean_Emulator/configs/samudra_om4_v1_attention/model.yaml) as the base example and [train.yaml](/scratch/ag11542/ocean/Ocean_Emulator/Ocean_Emulator/configs/samudra_om4_v1_attention/train.yaml) to launch training.

		BottleneckBlockType = Literal["attention", "transformer", "maxvit"]


		class AttentionBlockConfig(BaseConfig):

Uh oh!

Conversation

amogh-gulati commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

alxmrs commented Mar 9, 2026

Uh oh!

amogh-gulati commented Mar 9, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Mar 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alxmrs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

amogh-gulati Apr 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alxmrs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

amogh-gulati commented Mar 2, 2026 •

edited

Loading

amogh-gulati Apr 9, 2026 •

edited

Loading