Multi-Stream Attention Residuals

Attention Residuals (AttnRes) enables content-dependent retrieval over network depth, which can be viewed as single-stream softmax attention along the depth axis. Inspired by the success of multi-head attention along the sequence axis, we introduce Multi-Stream Attention Residuals (MSAR), an extension of AttnRes that generalizes this mechanism to multi-head depth attention, yielding substantial performance gains with minimal additional overhead. MSAR draws stream diversity from sequence context through multi-scale causal convolutions, allowing different heads to retrieve historical representations through distinct contextual views while preserving the integrity of the hidden representation. The retrieved streams are combined through lightweight adaptive mixing to generate the single-stream layer input. By retaining single-stream sources and expanding them on demand, MSAR supports multi-stream computation without requiring persistent multi-stream storage or communication. Across all 12 downstream benchmarks at the 18B and 28B MoE scales, MSAR consistently outperforms prior residual variants, including standard residual connections, manifold-constrained Hyper-Connections (mHC), and AttnRes. Scaling-law experiments yield an iso-loss compute speedup of 1.58times over the standard residual baseline, compared with 1.35times for AttnRes and 1.30times for mHC. These results make MSAR a promising residual architecture for large-scale pre-training.

한 줄 요약

MSAR은 다중 스트림 어텐션 잔여 연결을 통해 28B MoE 모델에서 1.58×의 계산 가속화를 달성한 새로운 잔여 아키텍처이다.

핵심 기여도

MSAR은 AttnRes를 다중 헤드 깊이 어텐션으로 확장하여 28B MoE에서 1.58×의 iso-loss compute speedup을 달성.
다중 스트림 생성은 multi-scale causal convolution을 통해 이루어지며, 각 스트림이 독립적인 깊이 어텐션을 수행.
MSAR은 기존 잔여 연결 대비 0.5%의 FLOPs 증가와 0.3%의 활성 파라미터 증가만으로 성능 향상.
18B 및 28B MoE에서 12개의 downstream task에서 기존 잔여 연결 방식을 모두 상회.

핵심 아이디어

기존 잔여 연결 방식은 깊이 방향에서 정보가 희석되는 문제를 해결하기 위해 AttnRes가 단일 스트림 softmax 어텐션을 도입했다. 그러나 이는 단일 헤드로 제한되어 표현력이 제약된다. MSAR은 시퀀스 축에서 다중 스트림을 생성하는 multi-scale causal convolution을 도입하여, 각 스트림이 독립적인 깊이 어텐션을 수행하도록 한다. 이는 `m`개의 스트림이 각각 다른 컨텍스트 스케일을 반영하면서도, 전체 잠재 표현의 일관성을 유지한다. 이후, `ℋ_pre` 스타일의 가벼운 adaptive mixing을 통해 단일 스트림 입력을 생성한다. 이 방식은 다중 헤드 어텐션의 표현력을 깊이 방향으로 이식하면서도, 저장 및 통신 비용을 최소화한다.

기술적 접근법

**MSAR 구성 요소**:
**Multi-scale causal convolutions**: `m`개의 스트림 생성.
**Independent depth-wise softmax attention**: 각 스트림이 독립적으로 이전 레이어의 표현을 검색.
**ℋ_pre-style adaptive mixing**: 가벼운 스칼라 믹싱을 통해 복수 스트림을 단일 입력으로 통합.
**계산 효율성**:
MSAR은 28B MoE에서 FLOPs 0.5%, 활성 파라미터 0.3% 증가.
스트림은 forward 시 causal convolution으로 생성되며, 메모리에 저장되지 않음.
`block mode`에서 `O(Nd)` 저장 비용으로 AttnRes와 동일한 효율성 유지.
**하이퍼파라미터**:
`m = 4` 스트림이 최적 성능, `S = 4` 블록 크기로 8개 블록 유지 시 대부분의 이점 회복.

주요 결과

**18B MoE 모델에서 학습 손실**: MSAR은 1.750 (Vanilla: 1.799, mHC: 1.776, AttnRes: 1.769).
**Compute speedup**: MSAR은 1.58× (Vanilla 대비), AttnRes: 1.35×, mHC: 1.30×.
**Downstream 성능**: 12개의 downstream task에서 MSAR이 기존 잔여 연결 방식을 모두 상회.
**Ablation study**: `m = 2` 스트림으로도 대부분의 이점 달성, `m = 4` 시 최대 성능.
**10B MoE 실험**: MSAR은 Pile loss 1.847 (Vanilla: 1.898, mHC: 1.871, AttnRes: 1.867).

의의 및 한계

MSAR은 잔여 연결의 표현력과 계산 효율성을 동시에 향상시킨 새로운 아키텍처로, 대규모 사전 학습에서 유용하다. 특히, `multi-stream` 접근법은 기존 잔여 연결 방식과는 다른 메커니즘으로, `cross-layer access`와 `multi-state representation`의 이점을 결합한다. 그러나 MSAR은 `m = 4` 이상의 스트림에서는 성능 개선이 미미하며, 스트림 수 증가에 따른 복잡도 증가를 고려해야 한다. 또한, `block mode`에서 `N ≈ 8` 블록 유지가 최적임을 보여주며, 더 많은 블록 분할은 이점을 감소시킨다. 이는 MSAR이 적절한 블록 크기와 스트림 수를 선택하는 데 주의가 필요함을 시사한다.

실용적 활용

MSAR은 대규모 언어 모델의 사전 학습에서 계산 효율성과 성능을 동시에 향상시키는 잔여 연결 방식으로, 특히 MoE 아키텍처와 결합하여 활용할 수 있다. `multi-scale causal convolution`을 통해 생성된 스트림은 다양한 컨텍스트 스케일을 반영하므로, 시퀀스 기반 작업(예: 추론, 코드 생성)에 유리하다. 또한, `block mode`와 `two-phase computation`을 지원하여 대규모 모델의 메모리 및 통신 비용을 최소화할 수 있어, 클라우드 기반 대형 언어 모델 학습에 적합하다.