NVIDIA / NVIDIA/Megatron-LM

Feature Request: Attention Residuals

Open
#4,016 9 comments 1 reaction 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

## Summary
Request to add support for [Attention Residuals (AttnRes)](https://arxiv.org/abs/2603.15031) to replace fixed-weight residual connections with learned, content-aware cross-layer attention, improving scaling efficiency by ~1.25x.

## Motivation
Standard residual connections accumulate all layer outputs with fixed unit weights, causing *PreNorm dilution* — as depth increases, individual layer contributions are diminished and hidden-state magnitudes grow unbounded.

AttnRes addresses this by computing selective softmax-attention over preceding layer outputs using a learned pseudo-query per layer, giving each layer content-aware access to earlier representations. The memory-efficient Block AttnRes variant partitions layers into ~8 blocks, applying attention only at block boundaries (O(Nd) vs O(Ld)), and matches baseline performance trained with 25% more compute. See [paper](https://arxiv.org/abs/2603.15031) for full results.

## Requested Features
1. **Block Attention Residuals** - Learnable cross-layer attention mechanism with
configurable block partitioning, replacing fixed residual connections in `TransformerLayer`
2. **MoE Testing** - Validation and integration with existing MoE architectures (paper uses Kimi Linear 48B)

## References
- [Attention Residuals Paper (Kimi/MoonshotAI)](https://arxiv.org/abs/2603.15031)
- [GitHub: MoonshotAI/Attention-Residuals](https://github.com/MoonshotAI/Attention-Residuals)
- Related: #2890

Contributor guide

Open the contributing guide

Research direction

Start by reading the TransformerLayer implementation, the linked Attention Residuals paper and reference repository, then inspect related issue #2890 for project context. Done means a configurable Block Attention Residuals integration is added to the residual path and validated with existing MoE architectures.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.