Kimi K3에서 채택된 Attention Residuals에 대해 정리해 보았다. 표준적인 잔차 연결 (Residual Connection)에 비해 고안된 점이 많다는 것을 잘 알 수 있다.

인용: https://arxiv.org/pdf/2607.24653

표준적인 잔차 연결의 과제

표준적인 잔차 연결의 경우 $l$번째 층에서

$d$는 벡터의 차원 크기

• 선행하는 모든 층의 출력이 하나의 $d$차원 벡터로 일률적으로 더해진다.
• 층 $l$이 과거 정보에 액세스하는 경로는 이 하나뿐이다. 즉, 깊이 방향의 모든 이력이 단일한 상태로 압축되어 있다.
  RNN의 은닉 상태 (Hidden State)가 과거의 모든 시점을 하나의 벡터에 채워 넣는 것과 동일한 병목 현상

잔차 연결 자체는 모델의 개량에 있어 매우 혁명적이라고 생각하므로, 여러 기사를 확인해 보시기 바란다.

Attention Residuals (Full AttnRes)에 관하여

학습 파라미터 $w$를 도입하여, 어텐션 (Attention) 메커니즘을 참고하여 Attention Residuals가 고안되었다. 다만 시퀀스 방향이 아니라 깊이 방향으로 적용되었다는 점이 핵심이며, 일반적인 어텐션과 비교하며 생각하면 이해하기 쉽다.

층의 단위에 관하여

층의 단위는 다음과 같다. 이후에 나오는 $h$는 층 단위의 값이다.

의사 쿼리 (Pseudo-query)와 학습 파라미터

$w_l$은 학습 가능한 파라미터 (Learnable Parameter) 그 자체이다.

그림의 $w$ 부분에 해당한다.

$q$는 의사 쿼리 (pseudo-query)이다.

일반적인 잔차 연결에 대해 학습 파라미터를 도입했다는 것이 큰 특징이다.

($k, v$ 부분에 관하여)

여기서 $f_i$는 $i$층으로부터의 출력 ($i$가 1인 경우는 임베딩 (Embedding) 결과인 $h_1$)

• $f(h)$는 누군가에게 더해져 사라지는 것이 아니라, 후속하는 모든 층으로부터 참조될 가능성이 있기 때문에 계속 유지해야 한다.

일반적인 어텐션과 비교했을 때:

• 루트(root)의 분모가 필요 없는 이유는 $w$가 생(raw) 파라미터이므로 경사 하강법 (Gradient Descent)이 노름 (Norm)을 적절한 크기로 조정할 수 있기 때문이다.
• RMSNorm에 의해 값의 상대치가 유지되면서 벡터 전체의 길이는 고정된다.

소프트맥스 커널 (Softmax Kernel)

Attention Residuals의 경우

일반적인 잔차 연결 (Residual Connection)과 비교했을 때, 상속을 하지 않는다는 것을 알 수 있다.

또한, 표준적인 잔차를 전개했을 경우 다음과 같은 식이 된다.

경로의 구조는 동일하고 계수가 $1
ightarrow \alpha$로 바뀌었을 뿐이라는 이해로도 문제없다.

K3가 실제로 채택한 것은 블록 입도 (Block Granularity)

Full AttnRes는 개념적인 정식화이며, 구현은 Block AttnRes이다. 층 단위가 아니라 블록 단위 (93층을 12층 $\times$ 8블록)로 선택하며, 블록 내부는 기존 방식대로 일률적인 가산을 수행한다.

$O(Ld)
ightarrow O(Nd)$의 메모리 절감과 추론 시의 상태 유계화 (State Boundedness)가 목적이며, $N \approx 8$일 때 이득의 대부분을 회수할 수 있다는 실험 결과에 기반하고 있다.

요약

일반적인 어텐션 형식을 참고하면서, 아래와 같은 차이는 있지만 과거 정보를 이전보다 효율적으로 취득하는 것이 가능해졌다.

Discussion