Add attention residuals - #873
Merged
Merged
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds a new “Full Attention Residual” routing option that replaces additive residual accumulation with token-local, depth-wise attention over the embedding plus all prior sublayer outputs, integrating it into the GPT forward pass with config/CLI support and accompanying tests/docs.
Changes:
- Introduces
FullAttentionResidual(depth-wise routing mixer) and a new model path gated byattention_residual_variant=full. - Extends
Blockto expose “no additive skip” attention/MLP transforms required by the new routing. - Adds configuration/CLI knobs plus tests, documentation, and an exploration YAML for comparisons.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| variations/block_variations.py | Adds helper entrypoints to run attention/MLP transforms without additive skips (used by full residual routing). |
| variations/attention_residual_variations.py | Implements the new depth-wise attention residual mixer. |
| model.py | Wires the new residual routing into GPT init and forward pass. |
| gpt_conf.py | Adds config fields for selecting residual routing variant and epsilon. |
| train_args.py | Adds CLI flags for attention_residual_variant and attention_residual_eps. |
| tests/test_attention_residual.py | Adds unit/integration tests covering forward+backward and initial equal-weight mixing behavior. |
| documentation/Attention_Residuals.md | Documents the architecture, constraints, and how to enable it. |
| explorations/default_inf_attention_residual_comparison.yaml | Adds an experiment spec to compare standard vs full attention residuals across sweeps. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+147
to
+153
| 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}") |
Comment on lines
+479
to
+499
| 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 |
Comment on lines
+23
to
+30
| 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) |
| muon_exclude_substrings: | ||
| - [" "] | ||
|
|
||
| # CAPPED HS NORM AND HSNORM VARIATONIS |
gkielian
approved these changes
Aug 2, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces a new "Full Attention Residual" architecture as an alternative to the standard additive residuals in the Transformer model. It implements depth-wise attention over all previous sublayer outputs, provides configuration and CLI support, integrates the new residual routing into the model's forward pass, and adds tests and documentation.
Implementation of Full Attention Residuals:
FullAttentionResidualmodule invariations/attention_residual_variations.py, which mixes previous sublayer outputs using learned pseudo-queries and token-local depth-wise attention.variations/block_variations.pyto addattention_residual_attnandattention_residual_mlpmethods, enabling extraction of raw attention and MLP outputs without additive skips, as required by the new residual routing.Model integration and configuration:
model.pyto integrate the full residual routing: added logic to select and run the new residual path, updated the forward pass to use it when configured, and ensured correct storage and mixing of sublayer outputs. [1] [2] [3] [4] [5]gpt_conf.pyandtrain_args.pyto add configuration options forattention_residual_variantand its epsilon parameter, with CLI support. [1] [2]Testing and documentation:
tests/test_attention_residual.pywith tests for the new module and model integration, ensuring correct behavior and gradients.documentation/Attention_Residuals.mdto describe the new architecture, its implementation, and usage instructions.Experimentation support:
explorations/default_inf_attention_residual_comparison.yamlto enable systematic comparison between standard and full attention residuals across various hyperparameter sweeps.