Add attention residuals

documentation/Attention_Residuals.md

@@ -0,0 +1,24 @@
+# Attention Residuals
+
+Set `attention_residual_variant: full` (or pass
+`--attention_residual_variant full`) to replace the running residual sum with
+token-local attention over depth.
+
+For each Transformer block, the implementation:
+
+1. mixes the embedding and earlier sublayer outputs for the attention input;
+2. appends the raw self-attention output to depth memory;
+3. computes a separate mixture for the MLP input; and
+4. appends the raw MLP output.
+
+One final mixture is passed to `ln_f`. Each destination owns a learned
+pseudo-query. Queries are initialized to zero, so every mixture initially is an
+equal-weight average. Keys are parameter-free RMS-normalized source vectors,
+values are the raw vectors, and softmax is only over depth. Consequently this
+feature does not replace causal token-to-token self-attention.
+
+Full Attention Residuals store the embedding and all `2 * n_layer` sublayer
+outputs, and perform quadratic work in the number of sublayers. The current
+implementation supports sequential attention-then-MLP blocks without post-LN;
+the usual PreNorm configuration is supported. Use `standard` (the default) for
+the existing additive residual architecture and checkpoint compatibility.

explorations/default_inf_attention_residual_comparison.yaml

@@ -0,0 +1,223 @@
+# Compare standard additive residuals with Full Attention Residuals while
+# preserving the architecture and sweeps from default_inf.yaml.
+---
+
+named_static_groups:
+  # QK Norm
+  - named_group: "qk_norm"
+    use_qk_norm: [true]
+    use_qk_norm_scale: [true]
+
+  # Norm Type
+  - named_group: "peri_ln"
+    use_pre_ln: [true]
+    use_peri_ln: [true]
+    use_post_ln: [false]
+
+  # Position Embeddings
+  - named_group: "rotary"
+    use_rotary_embeddings: [true]
+    use_abs_pos_embeddings: [false]
+
+  # MLP Activation
+  - named_group: "squared_relu"
+    activation_variant: ["squared_relu"]
+
+  # Attention softmax variants
+  - named_group: "relu2max"
+    softmax_variant_attn: ["relu2max"]
+
+  - named_group: "softmax"
+    softmax_variant_attn: ["softmax"]
+
+  # Infinite Attention
+  - named_group: "infinite"
+    attention_variant: ["infinite"]
+    use_concat_heads: [true]
+
+  # Head Dimension
+  - named_group: "hd_100"
+    n_qk_head_dim: [100]
+    n_v_head_dim: [100]
+
+  - named_group: "hd_150"
+    n_qk_head_dim: [150]
+    n_v_head_dim: [150]
+
+  - named_group: "hd_200"
+    n_qk_head_dim: [200]
+    n_v_head_dim: [200]
+
+  # MQA
+  - named_group: "mqa"
+    n_kv_group: [1]
+
+  # Residual architecture comparison
+  - named_group: "standard_residual"
+    attention_residual_variant: ["standard"]
+
+  - named_group: "full_attention_residual"
+    attention_residual_variant: ["full"]
+
+named_variation_groups:
+  - named_group: "head_dimension"
+    named_group_alternates: ["hd_100", "hd_150", "hd_200"]
+
+  - named_group: "residual_architecture"
+    named_group_alternates: ["standard_residual", "full_attention_residual"]
+
+common_group:
+  dataset: ["minipile"]
+  eval_interval: [2500]
+  max_iters: [10000]
+  never_save_checkpoint: [true]
+  compile: [true]
+  log_rankme: [true]
+  log_areq: [true]
+
+
+optimizer: ["muon"]
+weight_decay: [0.0]
+muon_momentum: [0.95]
+muon_ns_steps: [5]
+muon_nesterov: [true]
+muon_min_ndim: [2]
+
+parameter_groups:
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "squared_relu"
+      - "relu2max"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+
+  - named_group_static:
+      - "qk_norm"
+      - "peri_ln"
+      - "rotary"
+      - "softmax"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "squared_relu"
+      - "relu2max"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+    muon_exclude_substrings:
+      - [" "]
+
+  - named_group_static:
+      - "qk_norm"
+      - "peri_ln"
+      - "rotary"
+      - "softmax"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+    muon_exclude_substrings:
+      - [" "]
+
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "softmax"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+    muon_exclude_substrings:
+      - [" "]
+
+  # CAPPED HS NORM AND HSNORM VARIATONIS
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "squared_relu"
+      - "relu2max"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+    muon_exclude_substrings:
+      - [" "]
+    norm_variant_output: ["cappedhyperspherenorm", "rmsnorm"]
+
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "squared_relu"
+      - "relu2max"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+    muon_exclude_substrings:
+      - [" "]
+    norm_variant_output: ["cappedhyperspherenorm", "rmsnorm"]
+    norm_variant_wte: ["rmsnorm", "hyperspherenorm"]
+
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "squared_relu"
+      - "relu2max"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+    muon_exclude_substrings:
+      - [" "]
+    norm_variant_output: ["cappedhyperspherenorm", "rmsnorm", "hyperspherenorm"]
+    norm_variant_wte: ["rmsnorm", "hyperspherenorm"]
+    norm_variant_attn: ["cappedhyperspherenorm"]
+    use_peri_ln: [true, false]
+
+  # CAPPED HS NORM AND HSNORM VARIATONS WITHOUT THE exclude substrings
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "squared_relu"
+      - "relu2max"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+    norm_variant_output: ["cappedhyperspherenorm", "rmsnorm"]
+
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "squared_relu"
+      - "relu2max"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+    norm_variant_output: ["cappedhyperspherenorm", "rmsnorm"]
+    norm_variant_wte: ["rmsnorm", "hyperspherenorm"]
+
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "squared_relu"
+      - "relu2max"
+      - "infinite"
+      - "hd_100"
+    named_group_variations:
+      - "residual_architecture"
+    norm_variant_output: ["cappedhyperspherenorm", "rmsnorm", "hyperspherenorm"]
+    norm_variant_wte: ["rmsnorm", "hyperspherenorm"]
+    norm_variant_attn: ["cappedhyperspherenorm"]
+    use_peri_ln: [true, false]
+

gpt_conf.py

@@ -107,6 +107,10 @@ class GPTConfig:
     ln_f_input_mixer_variant: str = "linear"
     ln_f_mixer_top_k: int = 2
 
+    # Depth-wise residual routing. "full" stores every attention/MLP output.
+    attention_residual_variant: str = "standard"
+    attention_residual_eps: float = 1e-6
+
     # Attention Variation Specific
 
     ## Flash Lobo

model.py

@@ -47,6 +47,7 @@
 
 from shared_param_utils import SharedParamGroupCreator
 from variations.block_variations import Block
+from variations.attention_residual_variations import FullAttentionResidual
 
 class GPT(nn.Module):
 
@@ -143,6 +144,13 @@ def __init__(self, config):
 
         self.transformer['drop'] = nn.Dropout(config.dropout)
         self.transformer['h'] = nn.ModuleList([Block(config, mlp=shared_mlp_array[i], attn=shared_attn_array[i]) for i in range(config.n_layer)])
+        self.attention_residual_variant = config.attention_residual_variant
+        if self.attention_residual_variant == "full":
+            self.attention_residual = FullAttentionResidual(
+                2 * config.n_layer + 1, config.n_embd, config.attention_residual_eps
+            )
+        elif self.attention_residual_variant != "standard":
+            raise ValueError(f"unknown attention_residual_variant: {self.attention_residual_variant}")
         self.transformer['ln_f'] = norm_dictionary[config.norm_variant_output](config)
 
         # Optional post-embedding normalizations
@@ -256,6 +264,19 @@ def compute_lm_head_logits(self, x, lm_head_module):
         weight = self.apply_lm_head_norm(lm_head_module.weight)
         return F.linear(x, weight, lm_head_module.bias)
 
+    def _forward_full_attention_residual(self, x, iter_num):
+        """Run blocks while retaining each sublayer output as a depth source."""
+        sources = [x]
+        destination = 0
+        for block in self.transformer.h:
+            attn_input = self.attention_residual(sources, destination)
+            sources.append(block.attention_residual_attn(attn_input, iter_num))
+            destination += 1
+            mlp_input = self.attention_residual(sources, destination)
+            sources.append(block.attention_residual_mlp(mlp_input, iter_num))
+            destination += 1
+        return self.attention_residual(sources, destination)
+
     def _init_weights(self, module):
         """
         Custom weight initialization logic for GPT model.
@@ -441,7 +462,11 @@ def forward(self, idx, targets=None, iter_num=None, token_dict=None, target_dict
                 layer_outputs = [x]
 
             layer_idx = 1
-            for block in self.transformer.h:
+            blocks = self.transformer.h
+            if self.attention_residual_variant == "full":
+                x = self._forward_full_attention_residual(x, iter_num)
+                blocks = ()
+            for block in blocks:
                 x = block(x, iter_num)
 
                 # Steering logic
@@ -584,7 +609,11 @@ def forward(self, idx, targets=None, iter_num=None, token_dict=None, target_dict
                 layer_outputs = [x]
 
             layer_idx = 1
-            for block in self.transformer.h:
+            blocks = self.transformer.h
+            if self.attention_residual_variant == "full":
+                x = self._forward_full_attention_residual(x, iter_num)
+                blocks = ()
+            for block in blocks:
                 # Propagate tokens through layers
                 x = block(x, iter_num)
 

tests/test_attention_residual.py

@@ -0,0 +1,40 @@
+import torch
+
+from gpt_conf import GPTConfig
+from model import GPT
+from variations.attention_residual_variations import FullAttentionResidual
+
+
+def test_zero_queries_start_as_equal_weight_average():
+    mixer = FullAttentionResidual(n_destinations=1, n_embd=2)
+    sources = [
+        torch.tensor([[[1.0, 3.0]]]),
+        torch.tensor([[[5.0, 7.0]]]),
+    ]
+
+    result = mixer(sources, destination=0)
+
+    torch.testing.assert_close(result, torch.tensor([[[3.0, 5.0]]]))
+
+
+def test_full_attention_residual_model_forward_and_backward():
+    config = GPTConfig(
+        block_size=4,
+        vocab_size=32,
+        n_layer=2,
+        n_head=2,
+        n_kv_group=2,
+        n_embd=8,
+        dropout=0.0,
+        attention_residual_variant="full",
+    )
+    model = GPT(config)
+    tokens = torch.randint(0, config.vocab_size, (2, config.block_size))
+    targets = torch.randint(0, config.vocab_size, (2, config.block_size))
+
+    logits, loss = model(tokens, targets)
+    loss.backward()
+
+    assert logits.shape == (2, config.block_size, config.vocab_size)
+    assert model.attention_residual.queries.shape == (2 * config.n_layer + 1, config.n_embd)
+    assert model.attention_residual.queries.grad is not None

train_args.py

@@ -21,6 +21,13 @@ def parse_args():
     training_group = parser.add_argument_group('training_group')
     logging_group = parser.add_argument_group('logging_group')
 
+    model_group.add_argument(
+        '--attention_residual_variant', default='standard', choices=['standard', 'full'],
+        help='Residual stream implementation: ordinary addition or full depth-wise attention.',
+    )
+    model_group.add_argument('--attention_residual_eps', default=1e-6, type=float,
+                             help='RMSNorm epsilon used for Full Attention Residual routing keys.')
+
     # MLP Bias Configuration
     model_group.add_argument('--mlp_up_bias', default=None, action=argparse.BooleanOptionalAction, help='Whether to use bias in MLP up projections. If None, uses global bias setting.')
     model_group.add_argument('--mlp_down_bias', default=None, action=argparse.BooleanOptionalAction, help='Whether to use bias in MLP down projections. If None, uses global bias setting.')