From 1eec0ff5d813cfe5ee1b966ce1230f13c561b931 Mon Sep 17 00:00:00 2001
From: Kauna <16511995+klei22@users.noreply.github.com>
Date: Wed, 29 Jul 2026 14:43:50 -0700
Subject: [PATCH 1/3] feat(model): add full attention residuals

---
 documentation/Attention_Residuals.md        | 24 +++++++++++++
 gpt_conf.py                                 |  4 +++
 model.py                                    | 33 +++++++++++++++--
 tests/test_attention_residual.py            | 40 +++++++++++++++++++++
 train_args.py                               |  7 ++++
 variations/attention_residual_variations.py | 30 ++++++++++++++++
 variations/block_variations.py              | 22 ++++++++++++
 7 files changed, 158 insertions(+), 2 deletions(-)
 create mode 100644 documentation/Attention_Residuals.md
 create mode 100644 tests/test_attention_residual.py
 create mode 100644 variations/attention_residual_variations.py

diff --git a/documentation/Attention_Residuals.md b/documentation/Attention_Residuals.md
new file mode 100644
index 0000000000..b0546b8f0e
--- /dev/null
+++ b/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.
diff --git a/gpt_conf.py b/gpt_conf.py
index b8600a314e..056b7c7721 100644
--- a/gpt_conf.py
+++ b/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
diff --git a/model.py b/model.py
index e0e07239bf..1c4160d83c 100644
--- a/model.py
+++ b/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)
 
diff --git a/tests/test_attention_residual.py b/tests/test_attention_residual.py
new file mode 100644
index 0000000000..8d552cacb7
--- /dev/null
+++ b/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
diff --git a/train_args.py b/train_args.py
index 8e4497f7d9..456faab573 100644
--- a/train_args.py
+++ b/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.')
diff --git a/variations/attention_residual_variations.py b/variations/attention_residual_variations.py
new file mode 100644
index 0000000000..3ef7b41685
--- /dev/null
+++ b/variations/attention_residual_variations.py
@@ -0,0 +1,30 @@
+"""Depth-wise attention residuals.
+
+This module implements Full Attention Residuals: each Transformer sublayer gets
+an input selected from the embedding and all earlier sublayer outputs.  Routing
+is token-local (the softmax dimension is depth), so sequence mixing remains the
+responsibility of the normal self-attention module.
+"""
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+class FullAttentionResidual(nn.Module):
+    """Mix earlier sublayer outputs with zero-initialized pseudo-queries."""
+
+    def __init__(self, n_destinations: int, n_embd: int, eps: float = 1e-6):
+        super().__init__()
+        # Includes one destination for each attention/MLP and one for ln_f.
+        self.queries = nn.Parameter(torch.zeros(n_destinations, n_embd))
+        self.eps = eps
+
+    def forward(self, sources: list[torch.Tensor], destination: int) -> torch.Tensor:
+        if not sources:
+            raise ValueError("attention residuals require at least one source")
+        values = torch.stack(sources, dim=0)  # depth, batch, time, channels
+        keys = F.rms_norm(values, (values.size(-1),), eps=self.eps)
+        scores = torch.einsum("dbtc,c->dbt", keys, self.queries[destination])
+        weights = scores.softmax(dim=0)
+        return torch.einsum("dbt,dbtc->btc", weights, values)
diff --git a/variations/block_variations.py b/variations/block_variations.py
index f16a488b79..40e6ead7c4 100644
--- a/variations/block_variations.py
+++ b/variations/block_variations.py
@@ -476,6 +476,28 @@ def forward(self, x: torch.Tensor, iter_num: int):
             return checkpoint.checkpoint(self.block_forward, x, iter_num, use_reentrant=False)
         return self.block_forward(x, iter_num)
 
+    def attention_residual_attn(self, x: torch.Tensor, iter_num: int) -> torch.Tensor:
+        """Run only the attention transformation, without an additive skip."""
+        if self.use_parallel_mlp or self.use_edgellm_asic or self.use_post_ln_attn:
+            raise ValueError("Full Attention Residuals require a sequential block without post-attention norm")
+        out = self.attn(self.pre_ln_attn(x) if self.use_pre_ln_attn else x, iter_num)
+        if self.use_peri_ln_attn:
+            out = self.peri_ln_attn(out)
+        if self.attn_resid_scaler is not None:
+            out = self.attn_resid_scaler(out)
+        return out
+
+    def attention_residual_mlp(self, x: torch.Tensor, iter_num: int) -> torch.Tensor:
+        """Run only the MLP transformation, without an additive skip."""
+        if self.use_parallel_mlp or self.use_edgellm_asic or self.use_post_ln_mlp:
+            raise ValueError("Full Attention Residuals require a sequential block without post-MLP norm")
+        out = self.mlp(self.pre_ln_mlp(x) if self.use_pre_ln_mlp else x, iter_num)
+        if self.use_peri_ln_mlp:
+            out = self.peri_ln_mlp(out)
+        if self.mlp_resid_scaler is not None:
+            out = self.mlp_resid_scaler(out)
+        return out
+
     def _combine_resid(self, kind: str, x: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
         """Helper method to streamline forward block skip connections"""
         alpha = self.alpha_fns[kind](out)

From ffa16f6a67f779e6ed1de0ecb46bd54b57229da5 Mon Sep 17 00:00:00 2001
From: Kauna <16511995+klei22@users.noreply.github.com>
Date: Fri, 31 Jul 2026 16:07:30 -0700
Subject: [PATCH 2/3] feat(explorations): compare attention residuals

---
 ...ult_inf_attention_residual_comparison.yaml | 110 ++++++++++++++++++
 1 file changed, 110 insertions(+)
 create mode 100644 explorations/default_inf_attention_residual_comparison.yaml

diff --git a/explorations/default_inf_attention_residual_comparison.yaml b/explorations/default_inf_attention_residual_comparison.yaml
new file mode 100644
index 0000000000..aecdb3a8ac
--- /dev/null
+++ b/explorations/default_inf_attention_residual_comparison.yaml
@@ -0,0 +1,110 @@
+# 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]
+
+parameter_groups:
+  - named_group_static:
+      - "qk_norm"
+      - "peri_ln"
+      - "rotary"
+      - "squared_relu"
+      - "relu2max"
+      - "infinite"
+      - "mqa"
+    n_head:
+      range:
+        start: 1
+        end: 12
+        step: 1
+    named_group_variations:
+      - "head_dimension"
+      - "residual_architecture"
+
+  - named_group_static:
+      - "qk_norm"
+      - "peri_ln"
+      - "rotary"
+      - "softmax"
+      - "infinite"
+      - "mqa"
+    n_head:
+      range:
+        start: 1
+        end: 12
+        step: 1
+    named_group_variations:
+      - "head_dimension"
+      - "residual_architecture"

From a01283c3320db98a54edc76df6d8f6264f9583ee Mon Sep 17 00:00:00 2001
From: klei22 <kauna@avalonwest.tech>
Date: Sat, 1 Aug 2026 15:25:08 -0700
Subject: [PATCH 3/3] Update the residual comparison matrix

---
 ...ult_inf_attention_residual_comparison.yaml | 141 ++++++++++++++++--
 1 file changed, 127 insertions(+), 14 deletions(-)

diff --git a/explorations/default_inf_attention_residual_comparison.yaml b/explorations/default_inf_attention_residual_comparison.yaml
index aecdb3a8ac..9fd38ec3c8 100644
--- a/explorations/default_inf_attention_residual_comparison.yaml
+++ b/explorations/default_inf_attention_residual_comparison.yaml
@@ -75,23 +75,46 @@ common_group:
   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"
-      - "mqa"
-    n_head:
-      range:
-        start: 1
-        end: 12
-        step: 1
+      - "hd_100"
     named_group_variations:
-      - "head_dimension"
       - "residual_architecture"
+    muon_exclude_substrings:
+      - [" "]
 
   - named_group_static:
       - "qk_norm"
@@ -99,12 +122,102 @@ parameter_groups:
       - "rotary"
       - "softmax"
       - "infinite"
-      - "mqa"
-    n_head:
-      range:
-        start: 1
-        end: 12
-        step: 1
+      - "hd_100"
     named_group_variations:
-      - "head_dimension"
       - "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]
+