AI 요약
기존 트랜스포머 모델의 표준 잔차 연결(Residual Connections)은 모든 레이어의 출력을 동일한 가중치로 누적하여, 모델이 깊어질수록 각 레이어의 기여도가 희석되고 PreNorm 구조에서 은닉 상태(hidden-state)의 크기가 무한히 커지는 문제를 발생시킵니다. MoonshotAI는 이를 해결하기 위해 레이어별로 학습된 의사 쿼리(pseudo-query)를 사용하여 이전 레이어들의 출력을 선택적으로 집계하는 '어텐션 레지듀얼(AttnRes)' 기술을 공개했습니다. 이 방식은 고정된 가중치 대신 입력값에 의존적인 어텐션 메커니즘을 적용하여 모델이 필요한 시점에 필요한 정보를 더 효과적으로 참조할 수 있게 합니다. 또한, 메모리 효율성을 위해 전체 레이어를 N개의 블록으로 나누어 관리하는 '블록 AttnRes'를 제안함으로써 대규모 모델에서도 실용적으로 적용할 수 있는 기반을 마련했습니다.
핵심 인사이트
- 동적 가중치 할당: 표준 잔차 연결의 고정 단위 가중치 대신, $w_l \in \mathbb{R}^d$ 형태의 학습된 의사 쿼리를 통해 이전 레이어 표현들의 소프트맥스 어텐션 가중치를 계산합니다.
- 메모리 최적화: Full AttnRes의 $O(Ld)$ 메모리 복잡도 문제를 해결하기 위해 레이어를 N개 블록으로 그룹화하는 Block AttnRes를 도입했으며, 약 8개의 블록만으로도 대부분의 성능 향상을 회수합니다.
- PreNorm 안정화: PreNorm 구조에서 발생하는 은닉 상태의 무분별한 성장을 제어하여 심층 신경망의 학습 안정성을 높였습니다.
- 드롭인 교체(Drop-in Replacement): 기존 트랜스포머 아키텍처의 잔차 연결 부분을 최소한의 오버헤드로 즉시 대체할 수 있는 구조를 제공합니다.
주요 디테일
- 수학적 메커니즘: 각 레이어 $l$의 출력 $h_l$은 이전 모든 출력 $v_i$에 대해 $\alpha_{i \to l}$ 가중치를 곱한 합산으로 정의됩니다.
- 정규화 기법: 어텐션 로짓을 계산하기 전
RMSNorm을 적용하여 데이터의 스케일을 조정함으로써 학습의 일관성을 보장합니다. - 계산 효율성: PyTorch 구현 예시에서
torch.einsum을 활용해 텐서 연산을 최적화하고, 블록 단위로 연산을 분할하여 대규모 컨텍스트 처리를 용이하게 합니다. - 이중 적용 구조: Transformer 레이어 내부에서 셀프 어텐션(Self-Attention) 레이어 직전과 MLP 레이어 직전 두 지점에 각각 Block AttnRes가 적용됩니다.
- 코드 구조: 블록이 경계값(
block_size)에 도달하면 새로운 블록을 생성하고 이전 블록들을 리스트 형태로 관리하여 선택적 참조 범위를 설정합니다.
향후 전망
- 거대언어모델(LLM) 표준화: 모델의 깊이가 수백 층으로 깊어지는 최신 LLM 트렌드에서 학습 안정성을 위한 표준 아키텍처 구성 요소로 자리 잡을 가능성이 높습니다.
- 추론 효율성 연구: 어텐션 레지듀얼이 추론(Inference) 시 메모리 대역폭에 미치는 영향을 최적화하기 위한 추가적인 양자화나 하드웨어 가속 연구가 이어질 것으로 보입니다.
출처:hackernews