기술 요약: Depth-Recurrent Attention Mixtures (Dreamer)

문제 정의

대규모 언어 모델(LLM)의 전통적인 사고 사슬(Chain-of-Thought, CoT) 추론은 이산적인 자연어로 생각을 구체화하는 것에 의존합니다. 이러한 방식은 표현력을 제한하고 긴 시퀀스 길이로 인해 높은 계산 비용을 발생시킵니다. 반면, 깊이 재귀(Depth Recurrence, DR)—연속적인 공간에서 잠재적 추론을 가능하게 하기 위해 깊이에 따라 파라미터를 재사용하는 방식—는 유망한 대안을 제시하지만, 기존 구현체들은 두 가지 결정적인 병목 현상에 직면해 있습니다:

1. 레이어 크기 병목(Layer-Size Bottleneck): 밀집형(Dense) DR 모델을 확장하는 것은 비현실적입니다. 왜냐하면 모든 깊이 단계마다 전체 모델을 실행해야 하기 때문입니다.
2. 은닉 크기 병목(Hidden-Size Bottleneck): 추론을 시퀀스 차원에서 깊이 차원으로 밀어 넣는 것은 일정한 은닉 상태 크기의 한계를 악화시킵니다. 어텐션 메커니즘이 시퀀스 차원의 문제를 해결(Sequence Attention)한 것과 달리, 이전의 DR 연구들은 깊이 차원의 이 병목 현상을 간과했습니다.

또한, 기존 연구들은 연산량(FLOPs), 파라미터 수, 메모리 사용량을 엄격하게 맞춘 정교한 베이스라인이 부족하며, 종종 부분적으로 고정된 레이어 스택이나 밀집 코어에 의존하여 공정한 비교를 저해합니다.

방법론: Dreamer 프레임워크

저자들은 Depth-Recurrent Attention Mixtures의 모듈형 프레임 much인 Dreamer를 소개합니다. 이 아키텍처는 세 가지 직교하는 어텐션 차원을 하나의 반복 적용되는 레이어로 통합합니다:

1. 시퀀스 어텐션 (Sequence Attention, SA): 토큰 시퀀스 차원을 따라 작동하는 표준 셀프 어텐션입니다.
2. 깊이 어텐션 (Depth Attention, DA): 동일한 토큰의 이전 깊이들을 가로질러 어텐션을 적용하는 새로운 메커니즘입니다.
   • 구현: 시퀀스 차원을 배치(Batch) 차원으로 취급하고, 깊이를 시퀀스 차원으로 취급합니다.
   • 최적화: SA와 달리, DA의 키-값(KV) 캐시는 추론 중 각 토큰마다 버려지거나 덮어씌워질 수 있어, 시퀀스 길이에 관계없이 메모리 오버헤드를 일정하게 유지합니다.
   • 인코딩: 깊이를 위해 인코딩된 RoSE(Rotary Positional Embedding)를 사용하며, 최대 깊이에 대한 전방향 정보를 인코딩하기 위해 역방향 절반(reverse half)을 적용합니다.
3. 전문가 어텐션 (Expert Attention, EA): MLP 전문가(Experts) 상에서 작동하는 희소 혼합 전문가(Sparse Mixture of Experts, MoE)입니다.
   • DR을 위한 적응: 깊이 재귀를 처리하기 위해 투영 가중치(Projection weights)를 직접 공유하지 않습니다. 대신, 투영과 집계 과정을 가벼운 희소 MoEs로 전환하고, 훈련을 안정화하기 위해 결합된 어텐션 점수(Tied attention scores)를 사용합니다 지습니다.
   • 라우팅: 균형을 맞추기 위한 바이어스 벡터를 사용하는 Top-K 라우팅 전략을 사용하며(DeepSeek-V3와 유사), 계산 오버헤드 없이 그래디언트 흐름을 돕기 위해 스케일링된 공유형 상시 활성 전문가를 도입합니다.

아키텍처 설계:
모델은 이러한 어텐션 메커니즘을 순차적(DA $\to$ SA $\to$ EA)으로 적용하거나 병렬 변형으로 적용합니다. 저자들은 잔차 스트림(Residual stream)의 잔차 스케일을 관리하기 위해 RMSNorm을 활용합니다. 공정한 비교를 위해, 저자들은 이변량 좌표 하강법(Bivariate coordinate descent)을 사용하여 DR 모델과 고전적인 레이어드 베이스라인 간의 FLOPs, 파라미터 및 메모리 사용량을 엄격하게 일치시킵니다.

주요 기여

1. 통합 프레임워크: 시퀀스, 깊이, 전문가 어텐션을 결합한 모듈형 아키텍처를 도입하여, 효율적인 확장을 위해 스케일링 차원(파라미터, FLOPs, 깊이, 잠재 메모리)을 분리했습니다.
2. 병목 현상 완화:
   • **희소 전문가 어텐션(Sparse Expert Attention)**은 레이어 크기 병목을 완화하여 DR이 효율적으로 확장될 수 있게 합니다.
   • **깊이 어텐션(Depth Attention)**은 은닉 크기 병목을 완화하여 복잡한 잠재 추론을 위해 효과적으로 확장될 수 있게 합니다.
3. 엄격한 평가: 본 논문은 DR 변형 모델과 최첨단(SOTA) MoE 간의 긴밀하게 매칭된 비교(FLOP, 파라미터, 메모리)를 제공하며, 이는 기존 문헌에서 확인된 격차를 해소합니다.
4. 지식 사용 분석: 저자들은 깊이에 따른 지식 할당을 분석하여, 전문가 선택의 다양성과 재사용 패턴을 정량화합니다.

실험 결과

모델은 약 100B 토큰으로 학습되어 자연어 추론 벤치마크(GSM8K, MATH, MMLU-Math, MathQA) 및 언어 모델링 작업에 대해 평가되었습니다.

• 데이터 효율성: 깊이-재귀 모델(DR 및 DR+DA)은 FLOP, 파라미터 및 메모리가 매칭된 SOTA 베이스라인과 동일한 정확도에 도달하기 위해 2~8배 적은 학습 토큰을 필요로 했습니다.
• 성능 대 규모(Performance vs. Scale): 깊이 16(약 1B 파라미터)을 사용하는 모델이 모든 추론 벤치마크에서 깊이 32(약 2B 파라미터)를 가진 고전적 레이어드 베이스라인보다 우수한 성능을 보였습니다. 이는 더 뛰어난 성능을 내면서도 파라미터, FLOPs, 메모리를 약 2배 감소시킨 결과입니다.
• 깊이 어텐션의 영향: 순수 DR(DA 미포함)도 이점을 보였으나, 깊이 어텐션과 결합된 모델(DR+DA)이 가장 강력한 결과를 냈습니다. 특히 더 깊은 모델(깊이 32)에서 베이스라인 및 순수 DR보다 유의미하게 높은 성능을 기록했습니다.
• 전문가 다양성: 분석 결과, DR 모델은 SOTA MoE에 비해 2~11배 더 큰 전문가 선택 다양성을 보이는 것으로 나타났습니다. 전문가의 약 50%가 최소 7개의 서로 다른 깊이에서 사용되었으며, 이는 상당한 지식 재사용과 깊이 일반화를 나타냅니다.

의의 및 주장

본 논문은 Dreamer가 자연어 추론에서 긴밀하게 매칭된 SOTA Transformer를 지속적으로 능가하는 완전한 깊이-재귀 단일 레이어 아키텍처임을 주장합니다.

• 효율성: 프레임워크는 깊이 재귀를 통한 잠재 추론이 깊이 어텐션 및 희소 전문가와 결합될 때 우수한 데이터 및 파라미터 효율성을 제공함을 입증합니다.
• 표현력: 깊이 어텐션은 균일한 스킵 연결(Skip connections)을 넘어 깊이 간의 표적화된 정보 라우팅을 가능하게 함으로써 모델의 표현력을 질적으로 확장합니다.
• 모듈성: 지식 접근(시퀀스, 깊이, 전문가)을 통합된 어텐션 혼합으로 프레임링함으로써, 저자들은 모델 동작의 더 나은 이해, 제어 및 모듈형 확장성을 촉진하는 관점을 제안합니다.
• 확장성: 본 연구는 밀집 코어와 일정한 은 l 크기의 한계를 넘어 깊이를 확장해야 하는 결정적인 필요성을 다루며, 깊이를 "제1급 순차적 차원(First-class sequential dimension)"으로 취급해야 함을 시사합니다.

저자들은 향후 과제에 대해 신중한 태도를 유지하며, 동적 깊이(Dynamic depth)와 대안적인 위치 인코딩이 유망하지만, 학습된 깊이 범위를 넘어선 신뢰할 수 있는 일반화는 여전히 미개척 과제로 남아 있다고 언급했습니다. 또한, 잠재 추론이 결론에 도달하기 위해 모델이 취하는 특정 경로에 대한 해석 가능성 문제를 증가시킬 수 있음을 인정했습니다.