From 1255b9bac8ca3db24b6de024278ccde100e04dd0 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 11:53:53 -0700 Subject: [PATCH 01/10] Add full attention residual routing --- .../src/rg_nanogpt_one_head/model.py | 115 ++++++++++++++++-- 1 file changed, 105 insertions(+), 10 deletions(-) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py index af1c3447..c065db9f 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py @@ -18,12 +18,13 @@ class GPTConfig: dropout: float = 0.0 bias: bool = False tie_weights: bool = True + residual_mode: str = "standard" def __post_init__(self) -> None: - if self.n_layer != 1 or self.n_head != 1: + if self.n_layer < 1 or self.n_head != 1: raise ValueError( - "the reference architecture is fixed to one block and one " - "attention head" + "the reference architecture requires at least one block and " + "exactly one attention head" ) if self.n_embd % self.n_head != 0: raise ValueError("n_embd must be divisible by n_head") @@ -31,6 +32,10 @@ def __post_init__(self) -> None: raise ValueError("invalid GPT configuration") if not 0.0 <= self.dropout < 1.0: raise ValueError("dropout must be in [0, 1)") + if self.residual_mode not in {"standard", "full_attnres"}: + raise ValueError( + "residual_mode must be 'standard' or 'full_attnres'" + ) class LayerNorm(nn.Module): @@ -49,6 +54,44 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) +class RMSNorm(nn.Module): + """Parameter-free RMS normalization used only for AttnRes routing keys.""" + + def __init__(self, eps: float = 1e-6) -> None: + super().__init__() + self.eps = float(eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + scale = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + return x * scale.to(dtype=x.dtype) + + +class DepthAttentionRouter(nn.Module): + """Full Attention Residuals router over preceding depth states. + + Each routing point owns one learned pseudo-query vector. Previous residual + states are RMS-normalized to form keys while the unnormalized states remain + the values. Softmax is taken over depth, independently for every token. + """ + + def __init__(self, width: int) -> None: + super().__init__() + self.query = nn.Parameter(torch.zeros(width)) + self.norm = RMSNorm() + self.last_mean_weights: torch.Tensor | None = None + + def forward(self, states: list[torch.Tensor]) -> torch.Tensor: + if not states: + raise ValueError("AttnRes requires at least one residual state") + values = torch.stack(states, dim=0) # [depth, batch, time, width] + keys = self.norm(values) + logits = torch.einsum("d,nbtd->nbt", self.query, keys) + weights = logits.softmax(dim=0) + if not torch.jit.is_scripting(): + self.last_mean_weights = weights.detach().mean(dim=(1, 2)).cpu() + return torch.einsum("nbt,nbtd->btd", weights, values) + + class CausalSelfAttention(nn.Module): def __init__(self, cfg: GPTConfig) -> None: super().__init__() @@ -110,9 +153,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ).transpose(1, 2) dropout_p = self.dropout if self.training else 0.0 if q.device.type == "xla": - # Use core matmul/mask/softmax operations on TPU. This avoids - # depending on accelerator-specific SDPA kernel registration while - # preserving the same causal attention equation. y = self._xla_math_attention( q, k, @@ -150,15 +190,45 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class Block(nn.Module): def __init__(self, cfg: GPTConfig) -> None: super().__init__() + self.residual_mode = cfg.residual_mode self.ln1 = LayerNorm(cfg.n_embd, cfg.bias) self.attn = CausalSelfAttention(cfg) self.ln2 = LayerNorm(cfg.n_embd, cfg.bias) self.mlp = MLP(cfg) + self.attn_res_router = ( + DepthAttentionRouter(cfg.n_embd) + if cfg.residual_mode == "full_attnres" + else None + ) + self.mlp_res_router = ( + DepthAttentionRouter(cfg.n_embd) + if cfg.residual_mode == "full_attnres" + else None + ) def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.residual_mode != "standard": + states = self.forward_attnres([x]) + return states[-1] x = x + self.attn(self.ln1(x)) return x + self.mlp(self.ln2(x)) + def forward_attnres( + self, + states: list[torch.Tensor], + ) -> list[torch.Tensor]: + if self.attn_res_router is None or self.mlp_res_router is None: + raise RuntimeError("forward_attnres called for a standard block") + + attn_input = self.attn_res_router(states) + attn_state = attn_input + self.attn(self.ln1(attn_input)) + states.append(attn_state) + + mlp_input = self.mlp_res_router(states) + mlp_state = mlp_input + self.mlp(self.ln2(mlp_input)) + states.append(mlp_state) + return states + class GPT(nn.Module): def __init__(self, cfg: GPTConfig) -> None: @@ -202,8 +272,14 @@ def hidden_states(self, idx: torch.Tensor) -> torch.Tensor: x = self.drop( self.token_embedding(idx) + self.position_embedding(positions) ) - for block in self.blocks: - x = block(x) + if self.cfg.residual_mode == "full_attnres": + residual_states = [x] + for block in self.blocks: + residual_states = block.forward_attnres(residual_states) + x = residual_states[-1] + else: + for block in self.blocks: + x = block(x) return self.ln_f(x) def forward( @@ -221,7 +297,6 @@ def forward( return logits, loss def next_token_logits(self, idx: torch.Tensor) -> torch.Tensor: - # Apply the expensive vocabulary projection only to the final position. hidden = self.hidden_states(idx)[:, -1:, :] return self.lm_head(hidden) @@ -249,11 +324,31 @@ def generate_greedy( def parameter_count(self) -> int: return sum(parameter.numel() for parameter in self.parameters()) + def attention_residual_weights(self) -> dict[str, list[float]]: + """Return the latest mean depth-routing weights for diagnostics.""" + result: dict[str, list[float]] = {} + for block_index, block in enumerate(self.blocks): + for name, router in ( + ("ATTN", block.attn_res_router), + ("MLP", block.mlp_res_router), + ): + if router is None or router.last_mean_weights is None: + continue + result[f"L{block_index:02d}_{name}"] = [ + float(value) for value in router.last_mean_weights.tolist() + ] + return result + def transformer_matrix_items( model: GPT, ) -> list[tuple[str, str, int, torch.Tensor]]: - """Return the six transformer matrices used by WeightWatcher and Muon.""" + """Return Q/K/V/O and MLP matrices used by WeightWatcher and Muon. + + AttnRes pseudo-query vectors are intentionally excluded: they are 1-D + routing parameters, not transformer weight matrices, and therefore remain + in the auxiliary optimizer group rather than being folded into Muon/WW. + """ items: list[tuple[str, str, int, torch.Tensor]] = [] for block_index, block in enumerate(model.blocks): matrices = ( From abee03f44f20dedff3a6f99358354232b362116d Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 11:55:10 -0700 Subject: [PATCH 02/10] Add matched AttnRes long-horizon config --- .../configs/attnres_muon_10epochs.yaml | 106 ++++++++++++++++++ 1 file changed, 106 insertions(+) create mode 100644 baseline/nanogpt_one_head/configs/attnres_muon_10epochs.yaml diff --git a/baseline/nanogpt_one_head/configs/attnres_muon_10epochs.yaml b/baseline/nanogpt_one_head/configs/attnres_muon_10epochs.yaml new file mode 100644 index 00000000..6d95a371 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/attnres_muon_10epochs.yaml @@ -0,0 +1,106 @@ +protocol: + name: rg_nanogpt_one_head_attnres_muon_10epochs + version: 5 + description: Full Attention Residuals with ordinary Muon for ten corpus-equivalent epochs, exactly matched to the existing long-Muon baseline except for residual routing. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + residual_mode: full_attnres + +training: + seeds: [1337] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 10.0 + epoch_interval: 0.25 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.10 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 0.00006 + warmup_fraction: 0.01 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00003 + warmup_fraction: 0.05 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true From 178f5801292134d474df124973280848d2fcb2e7 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 11:55:25 -0700 Subject: [PATCH 03/10] Add conservative long-cosine Muon control --- .../configs/muon_10epochs_longcosine.yaml | 106 ++++++++++++++++++ 1 file changed, 106 insertions(+) create mode 100644 baseline/nanogpt_one_head/configs/muon_10epochs_longcosine.yaml diff --git a/baseline/nanogpt_one_head/configs/muon_10epochs_longcosine.yaml b/baseline/nanogpt_one_head/configs/muon_10epochs_longcosine.yaml new file mode 100644 index 00000000..6abed7fc --- /dev/null +++ b/baseline/nanogpt_one_head/configs/muon_10epochs_longcosine.yaml @@ -0,0 +1,106 @@ +protocol: + name: rg_nanogpt_one_head_muon_10epochs_longcosine + version: 5 + description: Standard residual one-head nanoGPT control using a conservative ten-epoch Muon cosine schedule for long-horizon convergence studies. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + residual_mode: standard + +training: + seeds: [1337] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 10.0 + epoch_interval: 0.25 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.02 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0004 + min_learning_rate: 0.00002 + warmup_fraction: 0.01 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW, conservative long cosine + family: muon + matrix_learning_rate: 0.01 + matrix_min_learning_rate: 0.0005 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00001 + warmup_fraction: 0.01 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true From 7cdc1b5f7c4467be852d3c9eddafcb75fdcfaced Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 11:55:39 -0700 Subject: [PATCH 04/10] Add conservative AttnRes long-cosine config --- .../attnres_muon_10epochs_longcosine.yaml | 106 ++++++++++++++++++ 1 file changed, 106 insertions(+) create mode 100644 baseline/nanogpt_one_head/configs/attnres_muon_10epochs_longcosine.yaml diff --git a/baseline/nanogpt_one_head/configs/attnres_muon_10epochs_longcosine.yaml b/baseline/nanogpt_one_head/configs/attnres_muon_10epochs_longcosine.yaml new file mode 100644 index 00000000..f96bad39 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/attnres_muon_10epochs_longcosine.yaml @@ -0,0 +1,106 @@ +protocol: + name: rg_nanogpt_one_head_attnres_muon_10epochs_longcosine + version: 5 + description: Full Attention Residuals one-head nanoGPT using the same conservative ten-epoch Muon cosine schedule as its standard-residual control. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + residual_mode: full_attnres + +training: + seeds: [1337] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 10.0 + epoch_interval: 0.25 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.02 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0004 + min_learning_rate: 0.00002 + warmup_fraction: 0.01 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW, conservative long cosine + family: muon + matrix_learning_rate: 0.01 + matrix_min_learning_rate: 0.0005 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00001 + warmup_fraction: 0.01 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true From 05a82571212dc8c98afe20bf1e12b1495331c352 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 11:56:04 -0700 Subject: [PATCH 05/10] Test attention residual baseline --- .../tests/test_attention_residuals.py | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 baseline/nanogpt_one_head/tests/test_attention_residuals.py diff --git a/baseline/nanogpt_one_head/tests/test_attention_residuals.py b/baseline/nanogpt_one_head/tests/test_attention_residuals.py new file mode 100644 index 00000000..33c4e3fb --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_attention_residuals.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import math + +import torch + +from rg_nanogpt_one_head.model import GPT, GPTConfig, transformer_matrix_items +from rg_nanogpt_one_head.optimizers import make_optimizer_handles + + +def _tiny_config(*, residual_mode: str, n_layer: int = 1) -> GPTConfig: + return GPTConfig( + vocab_size=97, + block_size=16, + n_layer=n_layer, + n_head=1, + n_embd=32, + dropout=0.0, + bias=False, + tie_weights=False, + residual_mode=residual_mode, + ) + + +def _muon_profile() -> dict: + return { + "family": "muon", + "matrix_learning_rate": 0.01, + "matrix_min_learning_rate": 0.0005, + "aux_learning_rate": 0.0003, + "aux_min_learning_rate": 0.00001, + "momentum": 0.95, + "nesterov": True, + "newton_schulz_steps": 3, + "muon_epsilon": 1.0e-7, + "matrix_weight_decay": 0.01, + "beta1": 0.90, + "beta2": 0.95, + "epsilon": 1.0e-8, + "aux_weight_decay": 0.01, + } + + +def test_full_attnres_forward_backward_is_finite() -> None: + torch.manual_seed(7) + model = GPT(_tiny_config(residual_mode="full_attnres")) + idx = torch.randint(0, model.cfg.vocab_size, (3, 12)) + targets = torch.randint(0, model.cfg.vocab_size, (3, 12)) + + logits, loss = model(idx, targets) + + assert logits.shape == (3, 12, model.cfg.vocab_size) + assert loss is not None and torch.isfinite(loss) + loss.backward() + + block = model.blocks[0] + assert block.attn_res_router is not None + assert block.mlp_res_router is not None + assert block.attn_res_router.query.grad is not None + assert block.mlp_res_router.query.grad is not None + assert torch.isfinite(block.attn_res_router.query.grad).all() + assert torch.isfinite(block.mlp_res_router.query.grad).all() + + +def test_one_block_routing_weights_have_expected_depth() -> None: + torch.manual_seed(11) + model = GPT(_tiny_config(residual_mode="full_attnres")) + idx = torch.randint(0, model.cfg.vocab_size, (2, 10)) + model(idx) + + weights = model.attention_residual_weights() + assert set(weights) == {"L00_ATTN", "L00_MLP"} + assert len(weights["L00_ATTN"]) == 1 + assert len(weights["L00_MLP"]) == 2 + assert math.isclose(weights["L00_ATTN"][0], 1.0, abs_tol=1e-6) + assert math.isclose(sum(weights["L00_MLP"]), 1.0, abs_tol=1e-6) + + +def test_transformer_matrices_remain_separate_from_attnres_queries() -> None: + model = GPT(_tiny_config(residual_mode="full_attnres")) + items = transformer_matrix_items(model) + + assert [item[1] for item in items] == [ + "W_Q", + "W_K", + "W_V", + "W_O", + "W_MLP_IN", + "W_MLP_OUT", + ] + assert len(items) == 6 + matrix_ids = {id(item[3]) for item in items} + assert id(model.blocks[0].attn_res_router.query) not in matrix_ids + assert id(model.blocks[0].mlp_res_router.query) not in matrix_ids + + +def test_muon_keeps_attnres_queries_in_auxiliary_optimizer() -> None: + model = GPT(_tiny_config(residual_mode="full_attnres")) + handles = make_optimizer_handles(model, _muon_profile()) + + assert [handle.role for handle in handles] == ["primary", "auxiliary"] + primary_ids = { + id(parameter) + for group in handles[0].optimizer.param_groups + for parameter in group["params"] + } + auxiliary_ids = { + id(parameter) + for group in handles[1].optimizer.param_groups + for parameter in group["params"] + } + + for _, _, _, matrix in transformer_matrix_items(model): + assert id(matrix) in primary_ids + for block in model.blocks: + assert id(block.attn_res_router.query) not in primary_ids + assert id(block.mlp_res_router.query) not in primary_ids + assert id(block.attn_res_router.query) in auxiliary_ids + assert id(block.mlp_res_router.query) in auxiliary_ids + + +def test_full_attnres_is_depth_ready_while_preserving_one_head() -> None: + torch.manual_seed(13) + model = GPT(_tiny_config(residual_mode="full_attnres", n_layer=3)) + idx = torch.randint(0, model.cfg.vocab_size, (2, 8)) + logits, _ = model(idx) + + assert logits.shape == (2, 8, model.cfg.vocab_size) + assert len(transformer_matrix_items(model)) == 18 + weights = model.attention_residual_weights() + assert len(weights["L00_ATTN"]) == 1 + assert len(weights["L00_MLP"]) == 2 + assert len(weights["L01_ATTN"]) == 3 + assert len(weights["L01_MLP"]) == 4 + assert len(weights["L02_ATTN"]) == 5 + assert len(weights["L02_MLP"]) == 6 From a00c7fe29e4155fd88b6583f45c5bf8adb93959d Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 11:56:31 -0700 Subject: [PATCH 06/10] Document Attention Residuals baseline --- .../nanogpt_one_head/ATTENTION_RESIDUALS.md | 123 ++++++++++++++++++ 1 file changed, 123 insertions(+) create mode 100644 baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md diff --git a/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md b/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md new file mode 100644 index 00000000..e4064bf3 --- /dev/null +++ b/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md @@ -0,0 +1,123 @@ +# Attention Residuals baseline + +This experiment adds **Full Attention Residuals (AttnRes)** as a controlled architectural baseline inside the existing one-head nanoGPT suite. + +## What changes + +Standard residual routing uses the immediately preceding state with a fixed additive path. Full AttnRes instead forms each sublayer input by content-dependent attention over all residual states available earlier in depth. Each routing point has a learned pseudo-query vector; prior residual states are RMS-normalized to form routing keys, while the original states remain the values. + +The implementation provides one router before attention and one router before the MLP. The existing causal self-attention computation itself is unchanged. + +## What does not change + +The controlled comparison keeps the existing long-Muon baseline fixed: + +- FineWeb-Edu sample and pinned revision +- GPT-2 tokenizer +- 80M training tokens, 1M validation tokens, 1M test tokens +- context length 256 +- one transformer block +- one attention head +- embedding width 128 +- zero dropout, bias disabled, tied token/output embeddings +- batch size 4 with 8 gradient-accumulation steps +- fixed evaluation probes and BLEU protocol +- WeightWatcher ERG/randomization settings +- separated `W_Q`, `W_K`, `W_V`, `W_O`, `W_MLP_IN`, and `W_MLP_OUT` matrices + +AttnRes pseudo-query vectors are one-dimensional parameter vectors. They are deliberately excluded from `transformer_matrix_items()`, so WeightWatcher continues to analyze the same six matrices per block. Under Muon, the six 2-D hidden matrices remain in the Muon group while the AttnRes queries enter the auxiliary AdamW group. + +## Why keep one block first? + +Full AttnRes becomes more expressive as depth increases. In a one-block transformer, the attention router initially has only the embedding residual state available, and the MLP router can select between the embedding state and the post-attention state. This is intentionally a conservative first experiment: it preserves the architecture, parameter scale, data exposure, and token budget of the current baseline and therefore isolates the effect of residual routing as cleanly as possible. + +The model implementation itself is depth-ready and keeps one head at every depth. A later multi-block study can test the larger routing advantage without changing the AttnRes implementation. + +## Experiment 1: exact matched long-Muon comparison + +Use the existing control: + +```bash +rg-onehead-train \ + --config configs/muon_10epochs.yaml \ + --optimizer muon \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-long-muon-standard/results \ + --device auto \ + --no-resume +``` + +Run AttnRes with every training hyperparameter unchanged: + +```bash +rg-onehead-train \ + --config configs/attnres_muon_10epochs.yaml \ + --optimizer muon \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-long-muon-attnres/results \ + --device auto \ + --no-resume +``` + +Use different results roots because run directories are keyed by optimizer and seed; the protocol fingerprint will correctly distinguish the model configurations but should not be forced to collide in one directory. + +## Experiment 2: conservative full-horizon cosine pair + +The existing ten-epoch Muon repair uses a one-epoch cosine schedule followed by nine epochs at the LR floor. That is the correct matched historical control, but it is not the only plausible schedule for studying long-horizon convergence. + +The new `*_longcosine.yaml` pair is a deliberately conservative candidate, not an empirically proven optimum. It uses the same schedule for both residual architectures: + +```text +training horizon: 10 epochs +LR schedule horizon: 10 epochs +Muon matrix peak LR: 0.0100 +Muon matrix floor LR: 0.0005 +auxiliary peak LR: 0.0003 +auxiliary floor LR: 0.00001 +warmup: 1% of the 10-epoch schedule +``` + +Run the standard-residual control: + +```bash +rg-onehead-train \ + --config configs/muon_10epochs_longcosine.yaml \ + --optimizer muon \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-longcosine-standard/results \ + --device auto \ + --no-resume +``` + +Run the matched AttnRes model: + +```bash +rg-onehead-train \ + --config configs/attnres_muon_10epochs_longcosine.yaml \ + --optimizer muon \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-longcosine-attnres/results \ + --device auto \ + --no-resume +``` + +Do not interpret the long-cosine pair as an AttnRes gain unless it is compared against its matching standard-residual long-cosine control. + +## Primary comparison metrics + +For convergence speed, compare at matched optimizer steps and matched tokens: + +- train, validation, and test loss +- validation and test next-token accuracy +- perplexity +- BLEU probe +- step/token count to fixed validation-loss thresholds +- best validation loss and the step at which it occurs + +For the RG/WeightWatcher analysis, retain the existing per-matrix metrics for all six matrices and compare trajectories of alpha, randomization distance, ERG quantities, and spectral diagnostics. AttnRes routing weights can additionally be inspected with `model.attention_residual_weights()` without altering the WeightWatcher matrix set. + +The key causal question is not whether the AttnRes run ends with a better number after ten epochs, but whether it reaches the same validation-loss or accuracy threshold in fewer matched tokens while preserving or improving out-of-sample behavior. From 1d3a39b926cae3e59e59c804e2eec1c472723feb Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 11:57:47 -0700 Subject: [PATCH 07/10] Correct Full AttnRes replacement semantics --- .../src/rg_nanogpt_one_head/model.py | 27 ++++++++++--------- 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py index c065db9f..824eb895 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py @@ -67,11 +67,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class DepthAttentionRouter(nn.Module): - """Full Attention Residuals router over preceding depth states. + """Full Attention Residuals router over preceding sublayer outputs. - Each routing point owns one learned pseudo-query vector. Previous residual - states are RMS-normalized to form keys while the unnormalized states remain - the values. Softmax is taken over depth, independently for every token. + Each routing point owns one learned pseudo-query vector. Previous outputs + are RMS-normalized to form keys while the unnormalized outputs remain the + values. Softmax is taken over depth, independently for every token. """ def __init__(self, width: int) -> None: @@ -82,7 +82,7 @@ def __init__(self, width: int) -> None: def forward(self, states: list[torch.Tensor]) -> torch.Tensor: if not states: - raise ValueError("AttnRes requires at least one residual state") + raise ValueError("AttnRes requires at least one preceding output") values = torch.stack(states, dim=0) # [depth, batch, time, width] keys = self.norm(values) logits = torch.einsum("d,nbtd->nbt", self.query, keys) @@ -220,13 +220,16 @@ def forward_attnres( if self.attn_res_router is None or self.mlp_res_router is None: raise RuntimeError("forward_attnres called for a standard block") + # Full AttnRes replaces additive residual accumulation. Each sublayer + # attends over the embedding and all preceding sublayer outputs to form + # its input h_l, and only f_l(Norm(h_l)) is appended as the next value. attn_input = self.attn_res_router(states) - attn_state = attn_input + self.attn(self.ln1(attn_input)) - states.append(attn_state) + attn_output = self.attn(self.ln1(attn_input)) + states.append(attn_output) mlp_input = self.mlp_res_router(states) - mlp_state = mlp_input + self.mlp(self.ln2(mlp_input)) - states.append(mlp_state) + mlp_output = self.mlp(self.ln2(mlp_input)) + states.append(mlp_output) return states @@ -273,10 +276,10 @@ def hidden_states(self, idx: torch.Tensor) -> torch.Tensor: self.token_embedding(idx) + self.position_embedding(positions) ) if self.cfg.residual_mode == "full_attnres": - residual_states = [x] + layer_outputs = [x] for block in self.blocks: - residual_states = block.forward_attnres(residual_states) - x = residual_states[-1] + layer_outputs = block.forward_attnres(layer_outputs) + x = layer_outputs[-1] else: for block in self.blocks: x = block(x) From 3647d4cba67e60ad6052490253e41502c42780a8 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 11:58:20 -0700 Subject: [PATCH 08/10] Clarify Full AttnRes sublayer semantics --- baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md b/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md index e4064bf3..034919a8 100644 --- a/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md +++ b/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md @@ -4,7 +4,7 @@ This experiment adds **Full Attention Residuals (AttnRes)** as a controlled arch ## What changes -Standard residual routing uses the immediately preceding state with a fixed additive path. Full AttnRes instead forms each sublayer input by content-dependent attention over all residual states available earlier in depth. Each routing point has a learned pseudo-query vector; prior residual states are RMS-normalized to form routing keys, while the original states remain the values. +Standard residual routing accumulates sublayer outputs through a fixed additive path. Full AttnRes replaces that accumulation: before each attention or MLP sublayer, it forms the sublayer input by content-dependent attention over the embedding and all preceding sublayer outputs. Each routing point has a learned pseudo-query vector; prior outputs are RMS-normalized to form routing keys, while the original outputs remain the values. The selected mixture is passed through the normal pre-norm sublayer, and only that sublayer output is appended as the next depth value. The implementation provides one router before attention and one router before the MLP. The existing causal self-attention computation itself is unchanged. @@ -25,11 +25,11 @@ The controlled comparison keeps the existing long-Muon baseline fixed: - WeightWatcher ERG/randomization settings - separated `W_Q`, `W_K`, `W_V`, `W_O`, `W_MLP_IN`, and `W_MLP_OUT` matrices -AttnRes pseudo-query vectors are one-dimensional parameter vectors. They are deliberately excluded from `transformer_matrix_items()`, so WeightWatcher continues to analyze the same six matrices per block. Under Muon, the six 2-D hidden matrices remain in the Muon group while the AttnRes queries enter the auxiliary AdamW group. +AttnRes pseudo-query vectors are 1-D parameter tensors of length `n_embd`. They are deliberately excluded from `transformer_matrix_items()`, so WeightWatcher continues to analyze the same six matrices per block. Under Muon, the six 2-D hidden matrices remain in the Muon group while the AttnRes queries enter the auxiliary AdamW group. ## Why keep one block first? -Full AttnRes becomes more expressive as depth increases. In a one-block transformer, the attention router initially has only the embedding residual state available, and the MLP router can select between the embedding state and the post-attention state. This is intentionally a conservative first experiment: it preserves the architecture, parameter scale, data exposure, and token budget of the current baseline and therefore isolates the effect of residual routing as cleanly as possible. +Full AttnRes becomes more expressive as depth increases. In a one-block transformer, the attention router initially has only the embedding output available, and the MLP router can select between the embedding and the attention-sublayer output. This is intentionally a conservative first experiment: it preserves the architecture, data exposure, token budget, and all six transformer matrices of the current baseline and therefore isolates the effect of residual routing as cleanly as possible. The only added trainable parameters are two length-128 pseudo-query vectors. The model implementation itself is depth-ready and keeps one head at every depth. A later multi-block study can test the larger routing advantage without changing the AttnRes implementation. From 052e8a9abe50ce0f416137cf5482e1ded24b45de Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 12:01:59 -0700 Subject: [PATCH 09/10] Add matched short AttnRes reference config --- .../configs/attnres_reference.yaml | 103 ++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 baseline/nanogpt_one_head/configs/attnres_reference.yaml diff --git a/baseline/nanogpt_one_head/configs/attnres_reference.yaml b/baseline/nanogpt_one_head/configs/attnres_reference.yaml new file mode 100644 index 00000000..6ba45cf3 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/attnres_reference.yaml @@ -0,0 +1,103 @@ +protocol: + name: rg_nanogpt_one_head_attnres_reference + version: 5 + description: Full Attention Residuals one-block, one-head nanoGPT matched exactly to the one-epoch FineWeb-Edu reference protocol except for residual routing. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + residual_mode: full_attnres + +training: + seeds: [1337, 2027, 4099] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 1.0 + epoch_interval: 0.125 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.10 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 0.00006 + warmup_fraction: 0.01 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00003 + warmup_fraction: 0.05 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true From 46b4bda5d806d00ed723e99eb7263ee5ed94e9e1 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 12 Aug 2026 12:02:51 -0700 Subject: [PATCH 10/10] Document short and long AttnRes CLI runs --- .../nanogpt_one_head/ATTENTION_RESIDUALS.md | 87 +++++++++++++++---- 1 file changed, 70 insertions(+), 17 deletions(-) diff --git a/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md b/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md index 034919a8..942ee0d7 100644 --- a/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md +++ b/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md @@ -10,7 +10,7 @@ The implementation provides one router before attention and one router before th ## What does not change -The controlled comparison keeps the existing long-Muon baseline fixed: +Every matched comparison keeps the existing one-head nanoGPT protocol fixed except for residual routing: - FineWeb-Edu sample and pinned revision - GPT-2 tokenizer @@ -27,15 +27,64 @@ The controlled comparison keeps the existing long-Muon baseline fixed: AttnRes pseudo-query vectors are 1-D parameter tensors of length `n_embd`. They are deliberately excluded from `transformer_matrix_items()`, so WeightWatcher continues to analyze the same six matrices per block. Under Muon, the six 2-D hidden matrices remain in the Muon group while the AttnRes queries enter the auxiliary AdamW group. -## Why keep one block first? +## Command-line setup -Full AttnRes becomes more expressive as depth increases. In a one-block transformer, the attention router initially has only the embedding output available, and the MLP router can select between the embedding and the attention-sublayer output. This is intentionally a conservative first experiment: it preserves the architecture, data exposure, token budget, and all six transformer matrices of the current baseline and therefore isolates the effect of residual routing as cleanly as possible. The only added trainable parameters are two length-128 pseudo-query vectors. +From `baseline/nanogpt_one_head`: -The model implementation itself is depth-ready and keeps one head at every depth. A later multi-block study can test the larger routing advantage without changing the AttnRes implementation. +```bash +python -m pip install -e . +export PYTORCH_ENABLE_MPS_FALLBACK=1 +``` + +Prepare the pinned corpus once and reuse it for every standard/AttnRes comparison: + +```bash +rg-onehead-prepare --config configs/reference.yaml +``` + +All scientific runs use the same command-line trainer: + +```text +rg-onehead-train +``` + +This preserves the existing checkpointing, restart, evaluation, WeightWatcher, MPS/CUDA/CPU/TPU device selection, and result layout. + +## Short run: exact one-epoch reference comparison -## Experiment 1: exact matched long-Muon comparison +This is the quickest clean comparison and uses all three canonical seeds by default. -Use the existing control: +Standard residual control: + +```bash +rg-onehead-train \ + --config configs/reference.yaml \ + --optimizer muon \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-short-standard/results \ + --device auto \ + --no-resume +``` + +AttnRes, with the same data, model dimensions, seeds, token budget, Muon hyperparameters, schedule, and probes: + +```bash +rg-onehead-train \ + --config configs/attnres_reference.yaml \ + --optimizer muon \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-short-attnres/results \ + --device auto \ + --no-resume +``` + +For a fast pilot before committing to all three seeds, append `--seeds 1337` to both commands. + +## Long run A: exact matched historical 10-epoch comparison + +This preserves the existing validated one-epoch Muon warmup/cosine horizon and then holds the LR floors through epoch 10. It is the correct historical comparison because it changes only residual routing. + +Standard residual control: ```bash rg-onehead-train \ @@ -48,7 +97,7 @@ rg-onehead-train \ --no-resume ``` -Run AttnRes with every training hyperparameter unchanged: +Matched AttnRes: ```bash rg-onehead-train \ @@ -61,13 +110,9 @@ rg-onehead-train \ --no-resume ``` -Use different results roots because run directories are keyed by optimizer and seed; the protocol fingerprint will correctly distinguish the model configurations but should not be forced to collide in one directory. - -## Experiment 2: conservative full-horizon cosine pair +## Long run B: convergence-oriented full-horizon cosine pair -The existing ten-epoch Muon repair uses a one-epoch cosine schedule followed by nine epochs at the LR floor. That is the correct matched historical control, but it is not the only plausible schedule for studying long-horizon convergence. - -The new `*_longcosine.yaml` pair is a deliberately conservative candidate, not an empirically proven optimum. It uses the same schedule for both residual architectures: +The historical 10-epoch protocol spends epochs 1--10 at the LR floor. For a long run intended to continue optimizing rather than primarily observe late-horizon behavior, the `*_longcosine.yaml` pair uses a conservative learning-rate schedule across all ten epochs: ```text training horizon: 10 epochs @@ -79,7 +124,9 @@ auxiliary floor LR: 0.00001 warmup: 1% of the 10-epoch schedule ``` -Run the standard-residual control: +The lower Muon peak relative to the one-epoch reference reduces long-horizon instability risk, while the nonzero floor prevents the optimization from becoming effectively frozen. This is a convergence-oriented candidate rather than a claim that its hyperparameters are globally optimal; actual convergence must be established from validation trajectories. + +Standard residual control: ```bash rg-onehead-train \ @@ -92,7 +139,7 @@ rg-onehead-train \ --no-resume ``` -Run the matched AttnRes model: +Matched AttnRes: ```bash rg-onehead-train \ @@ -105,7 +152,13 @@ rg-onehead-train \ --no-resume ``` -Do not interpret the long-cosine pair as an AttnRes gain unless it is compared against its matching standard-residual long-cosine control. +Do not attribute a gain to AttnRes unless it appears against the matching standard-residual config at the same token/step budget. + +## Why keep one block first? + +Full AttnRes becomes more expressive as depth increases. In a one-block transformer, the attention router initially has only the embedding output available, and the MLP router can select between the embedding and the attention-sublayer output. This is intentionally conservative: it preserves the architecture, data exposure, token budget, and all six transformer matrices of the current baseline and therefore isolates the effect of residual routing as cleanly as possible. The only added trainable parameters are two length-128 pseudo-query vectors. + +The model implementation itself is depth-ready and keeps one head at every depth. A later multi-block study can test the larger routing advantage without changing the AttnRes implementation. ## Primary comparison metrics @@ -120,4 +173,4 @@ For convergence speed, compare at matched optimizer steps and matched tokens: For the RG/WeightWatcher analysis, retain the existing per-matrix metrics for all six matrices and compare trajectories of alpha, randomization distance, ERG quantities, and spectral diagnostics. AttnRes routing weights can additionally be inspected with `model.attention_residual_weights()` without altering the WeightWatcher matrix set. -The key causal question is not whether the AttnRes run ends with a better number after ten epochs, but whether it reaches the same validation-loss or accuracy threshold in fewer matched tokens while preserving or improving out-of-sample behavior. +The key causal question is not whether the AttnRes run ends with a better number after a fixed horizon, but whether it reaches the same validation-loss or accuracy threshold in fewer matched tokens while preserving or improving out-of-sample behavior.