멀티-헤드 어텐션 잔차

초록

트랜스포머는 단일 가산 잔차 스트림(single additive residual stream)을 통해 깊이 방향으로 정보를 전파한다. 즉, 모든 하위 계층(sublayer)은 가장 최근 상태만 읽는다. 어텐션 잔차(attention residuals)는 각 하위 계층이 학습된 소프트맥스를 통해 깊이 이력에 주목함으로써 이 제약을 완화한다. 그러나 그 읽기는 전체 너비에 걸쳐 공유되는 단일 쿼리를 사용하므로, 모든 특징 부분공간(feature subspace)은 하나의 분포를 통해 깊이 이력을 읽어야 한다. 이러한 강제된 절충의 비용은 부분공간들이 어떤 계층을 읽을지에 대해 불일치하는 정도에 따라 증가하며, 그 불일치는 모델 너비에 따라 커진다. 우리는 멀티헤드 어텐션 잔차(Multi-Head Attention Residuals, MHAR)를 도입한다. 라우팅 쿼리를 H개의 부분공간별 헤드로 재구성하며, 각 헤드는 깊이 이력에 대한 자체 소프트맥스를 가진다. 읽기는 블록 대각(block-diagonal)이 되며, 재구성은 파라미터를 추가하지 않고 계산량도 무시할 수 있을 정도이다. H=1이면 어텐션 잔차와 정확히 동일해진다. 중복 제거되고 품질 필터링되었으며 STEM 및 코드 비중이 높은 Nemotron 기반 어닐 코퍼스(anneal corpus)에서 처음부터 학습했을 때, MHAR은 100M, 350M, 1B 규모에서 표준 트랜스포머 대비 검증 손실을 개선했다(-0.061, -0.149, -0.140). 모든 설정에서 네 가지 방법 중 최상의 결과를 달성했으며, 그 이득은 100M에서 더 큰 규모로 갈수록 증가했다. 헤드 수는 자유롭게 조절할 수 있는 손잡이가 아니라 실질적인 설계 축이다. 검증 손실은 H에 대해 U자형을 보이며, 모든 규모에서 H=4 또는 H=8에 평평한 최적점이 존재한다. 대규모 모델에는 H=8을 채택하며, 이 지점을 넘어선 과도한 분할(H=16)은 일관되게 이득의 일부를 상쇄한다. 학습된 쿼리에 대한 직접 프로빙은 학습된 부분공간 불일치가 근본적인 동인임을 확인한다. 융합된 Triton 라우팅 커널은 어텐션 잔차 학습 처리량을 기준 대비 0.2-0.5배에서 0.55-0.88배로 향상시키며, 최대 메모리는 기준과 거의 동일하게 유지한다. 델타 어텐션 잔차를 사용하는 항등성 보존 변환(identity-preserving conversion)은 8B 중간 학습을 지원하여 GSM8K에서 +3.2, GPQA에서 +3.1의 개선을 산출한다.

English

Transformers propagate information across depth through a single additive residual stream: every sublayer reads only the most recent state. Attention residuals relax this by letting each sublayer attend, through a learned softmax. However, that read uses a single query shared across the entire width, so every feature subspace must read the depth history through one distribution. The cost of this forced compromise grows with how much the subspaces disagree about which layers to read, and disagreement grows with model width. We introduce Multi-Head Attention Residuals (MHAR): the routing query is reshaped into H per-subspace heads, each with its own softmax over the depth history. The read becomes block-diagonal, the reshape adds zero parameters and negligible compute, and H = 1 recovers attention residuals exactly. Trained from scratch on a deduplicated Nemotron-based anneal corpus that is quality-filtered and STEM- and code-heavy, MHAR improves validation loss over a standard Transformer at 100M, 350M, and 1B (-0.061, -0.149, and -0.140). It achieves the best result among four methods in every setting, with the gain increasing from 100M to the larger scales. The head count is a real design axis rather than a free knob: validation loss is U-shaped with respect to H, with a flat optimum at H = 4 or H = 8 across scales. We adopt H = 8 for large-scale models; over-splitting beyond this point (H = 16) consistently gives back part of the gain. A direct probe of the trained queries confirms that learned subspace disagreement is the underlying driver. Fused Triton routing kernels increase attention-residual training throughput from 0.2-0.5x to 0.55-0.88x of the baseline while maintaining near-baseline peak memory. An identity-preserving conversion using delta attention residuals supports 8B mid-training, yielding improvements of +3.2 on GSM8K and +3.1 on GPQA.