깊이 혼합 어텐션: 대규모 언어 모델의 정보 희석 문제 해결을 위한 동적 계층 간 정보 추출 기법

모델이 깊어질수록 초기에 학습한 중요한 정보가 뒤로 갈수록 희미해지는 '정보 희석' 문제를 해결한다. 기존의 단순한 잔차 연결을 넘어, 필요한 시점에 과거 계층의 정보를 직접 꺼내 쓸 수 있게 하여 모델의 효율성과 성능을 동시에 개선했다.

Transformer는 계층을 쌓을수록 더 복잡한 추론이 가능해지지만, 실제로는 잔차 연결을 통해 정보가 계속 더해지면서 초기 계층의 특징이 희석되는 문제가 발생한다. 이는 깊은 층으로 갈수록 초기의 중요한 단서를 잃어버리게 만든다. MoDA는 이를 해결하기 위해 '세로 방향의 어텐션'을 도입한다. 기존 어텐션이 문장 내 다른 단어를 보는 가로 방향의 연산이라면, MoDA는 현재 단어 위치에서 이전 층들이 계산했던 결과물을 직접 들여다보는 세로 방향의 연산을 수행한다. 즉, 하나의 Softmax 연산으로 현재 층의 문맥 정보와 과거 층의 깊이 정보를 동시에 통합하여 처리한다. 결과적으로 모델은 매 층마다 지금 이 단어를 이해하기 위해 문장의 다른 단어를 볼 것인가, 아니면 이전 층에서 처리했던 이 단어의 핵심 정보를 다시 가져올 것인가를 데이터에 따라 동적으로 결정한다.

MoDA는 Transformer 블록의 '읽기-연산-쓰기' 과정을 재정의한다. 기존의 단순 잔차 연결 대신, 이전 모든 계층의 KV 쌍을 '깊이 메모리'로 간주하고 현재 계층의 Query가 이를 참조할 수 있도록 통합 Softmax 구조를 설계했다. 수학적으로는 현재 계층 l의 Query Q_l이 시퀀스 KV {K_l, V_l}과 깊이 KV {K_i, V_i}를 동시에 입력으로 받는다. [입력값 Q, K, V → 통합 Softmax 연산 → 출력값 O] 순으로 계산되며, 모든 어텐션 스코어는 하나의 Softmax 함수 내에서 정규화되어 시퀀스 정보와 깊이 정보가 동일한 표현 공간에서 융합된다. 하드웨어 효율성을 위해 Chunk-aware Depth-KV 레이아웃을 도입했다. 메모리 상에서 각 토큰의 깊이 방향 데이터를 연속적으로 배치하여 GPU의 Tensor Core 활용도를 극대화하고 HBM 대역폭 낭비를 최소화했다.

1.5B 파라미터 모델 실험에서 MoDA는 OLMo2 베이스라인 대비 10개 검증 벤치마크에서 평균 Perplexity를 0.2 낮추는 성과를 거두었다. 특히 C4 데이터셋에서 일관되게 낮은 손실 값을 기록하며 학습 안정성을 증명했다. 다운스트림 태스크 성능에서도 HellaSwag, WinoGrande, ARC-Challenge 등 10개 주요 평가 항목에서 평균 2.11%의 정확도 향상을 보였다. 연산 오버헤드는 FLOPs 기준 3.7%에 불과하여 효율적인 성능 확장이 가능함을 입증했다. 추가 분석을 통해 MoDA가 기존 Transformer의 고질적인 문제인 Attention Sink 현상을 완화하고, 정보가 풍부한 깊이 KV에 확률을 더 적절히 배분함을 확인했다.

MoDA 아키텍처는 기존의 Depth Residual과 Depth Dense의 장점을 결합한 형태다. Depth Dense가 O(L^2 D^2)의 파라미터 증가를 초래하는 반면, MoDA는 통합 Softmax를 통해 O(LD^2/G) 수준으로 파라미터 효율성을 높였다. 핵심 알고리즘은 하드웨어 친화적인 Fused Kernel로 구현되었다. Online Softmax 상태를 공유하며 시퀀스 어텐션과 깊이 어텐션을 한 번의 Forward Pass에서 처리한다. 특히 Group-aware Indexing을 통해 GQA 구조와의 호환성을 확보했다. 실험 결과 Post-norm 구조가 Pre-norm보다 MoDA와 결합했을 때 더 우수한 성능을 보였다. 이는 깊이 방향의 정보 흐름이 Post-norm의 신호 보존 특성과 더 잘 맞물리기 때문으로 분석된다.

대규모 산업용 학습을 위한 최종 단계의 최적화는 아직 과제로 남아 있으며, 모든 과거 계층의 KV를 캐싱하는 데 따르는 메모리 오버헤드가 깊이가 매우 깊어질 경우 병목이 될 수 있다.