본 논문은 최신 LLM에서 표준으로 사용되는 PreNorm 방식의 Residual connection이 모든 레이어 출력을 고정된 단위 가중치로 누적하여 은닉 상태(hidden-state)의 비제어적 성장을 야기하고, 각 레이어의 기여도를 점진적으로 희석시키는 문제점을 지적한다. 이러한 깊이 방향(depth-wise)의 정보 집합 방식이 고정되어 있다는 점에서, 시퀀스 차원에서 RNN이 직면했던 문제점과 유사한 "시간(time)과 깊이(depth)의 이중성"을 제안한다.

이 문제에 대한 해결책으로 Attention Residuals(AttnRes)를 제안한다. AttnRes는 고정된 누적 방식 $\left(h_l = \sum_{i} v_i\right)$을 각 레이어가 이전 레이어의 출력을 학습된 입력 의존적 가중치로 선택적으로 통합하는 Softmax Attention $\left(h_l = \sum_{i} \alpha_{i \to l} \cdot v_i\right)$으로 대체한다.

핵심 방법론은 다음과 같다.

1. Full Attention Residuals:
각 레이어 $l$의 입력 $h_l$은 이전 모든 레이어의 출력 $v_i$에 대한 Softmax Attention을 통해 계산된다.
$$h_l = \sum_{i=0}^{l-1} \alpha_{i \to l} \cdot v_i$$
여기서 $\alpha_{i \to l}$은 다음과 같이 정의되는 Attention weight이다.
$$\alpha_{i \to l} = \frac{\phi(q_l, k_i)}{\sum_{j=0}^{l-1} \phi(q_l, k_j)}$$
여기서 커널 함수 $\phi(q, k) = \exp \left( q^\top \text{RMSNorm}(k) \right)$를 사용한다. Query $q_l$은 각 레이어 $l$에 대한 학습 가능한 D-차원 벡터인 $w_l$이며, Key $k_i$ 및 Value $v_i$는 다음과 같이 정의된다.
$$v_i = \begin{cases} h_1 & i = 0 \\ f_i(h_i) & 1 \le i \le l-1 \end{cases}$$
$h_1$은 토큰 임베딩을 나타내며, $f_i(h_i)$는 레이어 $i$의 변환 출력을 나타낸다. RMSNorm의 사용은 큰 크기의 출력값이 Attention weight를 지배하는 것을 방지한다.
이 방식은 각 토큰에 대해 $O(L^2d)$의 연산량과 $O(Ld)$의 메모리를 요구한다. 이는 역전파를 위해 이미 유지되는 활성화(activation)와 메모리 오버랩되므로, 바닐라 트레이닝에서는 추가 메모리 오버헤드가 없다. 하지만 대규모 분산 트레이닝에서는 활성화 재계산 및 파이프라인 병렬화로 인해 $O(Ld)$의 통신 오버헤드가 발생한다.

2. Block Attention Residuals:
Full AttnRes의 $O(Ld)$ 메모리 및 통신 오버헤드를 줄이기 위해, Block AttnRes는 총 $L$개의 레이어를 $N$개의 블록으로 분할한다. 각 블록 내에서는 표준 Residual connection을 통해 레이어 출력을 단일 Representation으로 요약한다. 그리고 블록 간에는 $N$개의 블록 레벨 Representation에 대해서만 Full Attention을 적용한다.
블록 $n$의 레이어 인덱스 집합을 $B_n$이라고 할 때, 블록 표현 $b_n$은 다음과 같다.
$$b_n = \sum_{j \in B_n} f_j(h_j)$$
토큰 임베딩 $h_1$을 $b_0$으로 정의하여 항상 소스로 포함한다. 블록 $n$의 $i$번째 레이어에 대한 Value matrix는 다음과 같다.
$$V = \begin{cases} [b_0, b_1, \dots, b_{n-1}]^\top & \text{if } i = 1 \text{ (first layer of block } n) \\ [b_0, b_1, \dots, b_{n-1}, b_{i-1}^{(n)}]^\top & \text{if } i \ge 2 \text{ (subsequent layers)} \end{cases}$$
여기서 $b_{i-1}^{(n)}$은 블록 $n$ 내의 $i-1$번째 레이어까지의 부분합(partial sum)을 나타낸다. 이 방식은 메모리와 통신 오버헤드를 $O(Nd)$로 줄인다. $N$은 일반적으로 $L$보다 훨씬 작다 (논문에서는 $N \approx 8$이 대부분의 이점을 얻는다고 언급).

3. Infrastructure Design:

• Training (파이프라인 병렬화): AttnRes는 각 단계에서 누적된 모든 블록 Representation을 필요로 한다. naive하게 모든 히스토리를 전송하면 $O(C^2Nd)$의 통신 비용이 발생한다. 이를 최적화하기 위해, 'Cross-stage caching'을 사용한다. 이전 가상 단계에서 수신된 블록은 로컬 메모리에 캐시되어 재전송되지 않는다. 이로 인해 통신 비용이 $O(P^2Nd)$로 줄어들고, 통신과 연산의 오버랩이 가능해진다.
• Inference (추론):
  • Two-phase computation strategy:
    • Phase 1 (Parallel inter-block attention): 블록 내의 모든 $S=L/N$개 레이어에 대한 쿼리를 단일 배치(batch) 행렬 곱셈으로 동시에 계산한다. 이를 통해 $N$개의 캐시된 블록 Representation에 대한 메모리 접근 비용이 상각(amortize)되어, 각 레이어당 메모리 I/O를 $N/S$만큼 줄인다.
    • Phase 2 (Sequential intra-block attention + Online softmax merge): 각 레이어에 대해 순차적으로 intra-block attention을 계산하고, 온라인 Softmax를 통해 Phase 1의 결과와 병합한다. 이는 커널 퓨전(kernel fusion)을 가능하게 하여 I/O 오버헤드를 더욱 줄인다.
  • Memory-efficient prefilling: 긴 컨텍스트(long-context) 입력에 대한 prefilling 시 블록 Representation 저장을 위해 $N \cdot T \cdot d$ 메모리가 필요하다. 이를 줄이기 위해, 시퀀스 차원을 따라 Tensor Parallel 장치에 Representation을 샤딩(sharding)하여 각 장치에서 독립적으로 Phase 1을 실행한다. Phase 2의 온라인 Softmax 병합은 표준 TP all-reduce 통신 경로에 통합된다.

결론:
AttnRes는 PreNorm 희석 문제를 완화하고, 깊이에 따른 출력 크기를 제한하며, 기울기 분포를 더 균일하게 만든다. Scale law 실험을 통해 계산 예산 전반에 걸쳐 일관된 성능 향상을 보였으며, Block AttnRes는 baseline 대비 1.25배 더 많은 연산으로 훈련된 모델과 동등한 손실을 달성했다. Kimi Linear 아키텍처 (48B 파라미터)에 통합하여 1.4T 토큰으로 사전 훈련한 결과, 모든 downstream 태스크에서 baseline 대비 성능 향상을 입증했다.