깊이 방향 어텐션을 다중 스트림으로 확장한 MSAR

무엇인가

이 논문이 다루는 문제는 깊이 스케일링의 효율 저하다. PreNorm Transformer에서 잔차 스트림은 이전 계산을 단일 은닉 상태 궤적으로 누적하는데, 층이 깊어질수록 초기 층에서 형성된 특징이 희석되어 후속 층이 꺼내 쓰기 어려워진다. 논문은 표준 잔차 누적을 h_{l+1} = h_l + f_l(h_l)로 두고 이를 풀면 h_l = Σ v_i, 즉 모든 이전 소스가 고정 계수 1로 하나의 스트림에 합쳐지는 '고정 계수 깊이 혼합'이라고 규정한다. 이 관점에서 Hyper-Connections(mHC) 계열은 여러 개의 지속 스트림을 유지하는 순환형 선형 어텐션 스타일 깊이 메모리로, AttnRes는 깊이 축에 대한 단일 헤드 소프트맥스 어텐션으로 해석된다.

어떻게 동작하나

저자들이 지적하는 핵심 설계 난제는 단일 층 출력에서 다양성을 확보한 깊이 어텐션 헤드를 어떻게 만들 것인가다. AttnRes는 이미 은닉 상태를 채널 분할해 다중 헤드를 시도했지만 성능이 오히려 나빠졌는데, 깊이 방향 어텐션에서 검색되는 값이 은닉 상태 그 자체이므로 저차원 조각으로 자르면 표현의 완결성이 깨지기 때문이다. mHC의 H^post 방식으로 후보 스트림을 여러 개 만들면 각 후보가 같은 벡터의 스칼라 재가중에 불과하고, 정규화가 크기 차이를 지워 키가 거의 동일해져 다양성이 사라진다.

무엇과 다른가

MSAR의 해법은 다양성을 깊이가 아니라 시퀀스 축에서 끌어오는 것이다. 각 잔차 소스 v_i에 대해 항등 변환(r=0)과 커널 크기 {4, 8, 12}의 인과적 depthwise 1D 합성곱(r=1,2,3)을 적용해 m=4개의 전차원 스트림 v_i^(r)을 만든다. 이 스트림들은 서로 다른 국소 문맥 범위를 인코딩하므로, 각 스트림이 자체 학습된 의사 쿼리 w_l^(r)로 깊이 축 소프트맥스 검색을 독립 수행한다(스트림별로 가중치 합은 1). 검색된 m개의 상태 Z_l은 mHC 스타일의 경량 H^pre 믹서로 융합되는데, Flatten(Z_l)에 RMSNorm을 걸고 md×m 투영과 시그모이드를 통과시켜 입력 의존적 스칼라 가중치를 만들고 이를 스트림들의 가중합으로 층 입력 h_l을 구성한다. 즉 깊이 어텐션 프리미티브 자체를 다중 헤드로 넓힌 것이지, 지속적인 다중 스트림 순환을 추가한 것이 아니다.

어떻게 쓰나

시스템 측면의 핵심은 저장·통신 비용을 늘리지 않는다는 점이다. Block MSAR은 L개의 잔차 분기를 N개 블록으로 나누고 블록별 단일 스트림 요약 b_n만 유지한 뒤, 깊이 어텐션이 필요할 때 합성곱으로 스트림을 온디맨드 확장한다. 확장된 스트림은 곧바로 어텐션에 소비되어 다른 소비자가 없으므로 합성곱 확장과 스트림별 어텐션을 융합 커널로 처리해 중간값 쓰기를 없앨 수 있다. 그 결과 N개 소스에 대해 파이프라인 통신량과 프리필 활성값 저장량이 O(mNd)에서 O(Nd)로 줄어 AttnRes와 동일해진다. AttnRes의 단계 간 캐싱과 2단계 계산(블록 내 모든 타깃 쿼리를 완료된 소스에 배치 처리한 뒤, 각 층이 자기 부분합 기여분을 온라인 소프트맥스로 병합)도 그대로 호환된다. 28B MoE 기준 추가 부담은 학습 FLOPs 0.5%, 활성 파라미터 0.3%에 불과하다.

전제와 한계

실험은 DeepSeekMoE 스타일 Transformer, SwiGLU FFN, 18B·28B MoE 규모에서 진행됐고 기본 설정은 m=4 스트림, N≈8 블록이다. 12개 벤치마크(MMLU, MMLU-Pro, MMLU-Redux, ARC-C, BBH / HumanEval, LCBench, GSM8K, MATH / CMMLU, C-Eval, C3) 전부에서 MSAR가 Vanilla, mHC, AttnRes를 앞섰다. 18B에서 평균 45.7로 Vanilla 37.5(+8.2), mHC 41.7(+4.0), AttnRes 41.8(+3.9)였고, AttnRes 대비 MMLU +3.5, ARC-C +6.2, BBH +5.9였다. 28B에서는 평균 52.4로 Vanilla 44.6(+7.8), mHC 47.4(+5.0), AttnRes 49.4(+3.0)다. 18B 학습 손실은 MSAR 1.750, Vanilla 1.799, mHC 1.776, AttnRes 1.769였다. 스케일링 법칙 실험은 4개 모델 규모에서 L(C)=AC^{-β}+E(E=0.72)를 로그 공간 선형회귀로 적합했고, 최대 규모에서 손실을 Vanilla 대비 2.7% 줄였으며(mHC 1.6%, AttnRes 1.8%), 4개 규모 평균 iso-loss 연산 속도 향상은 1.58배로 AttnRes 1.35배, mHC 1.30배를 웃돈다.

10B MoE(15개 Transformer 층, 30개 잔차 서브층) Pile 손실 절제 실험도 설계 근거를 뒷받침한다. Vanilla 1.898, mHC 1.871, AttnRes 1.867에 대해 MSAR는 1.847이다. AttnRes의 이력 소스와 mHC의 다중 스트림 상태를 동시에 유지하는 하이브리드(AttnRes+mHC)는 1.861로 AttnRes 대비 0.006만 개선됐고, 같은 m=4·S=4 조건의 MSAR는 여기서 0.014를 더 낮춘다. 합성곱을 제거해 동일 소스를 4개 스트림으로 복제하고 쿼리만 독립적으로 둔 변형은 1.862로 나빠져, 쿼리 다양성만으로는 부족하고 소스 다양성이 중요함을 보인다. H^pre를 평균 집계로 바꾸면 1.851로 소폭 악화되는데, 적응적 융합이 추가 이득을 주면서도 균등 평균이 이득의 상당 부분을 유지한다는 뜻이다. 스트림 수는 m=2가 1.850, m=4가 1.847, m=8이 1.846으로 4 이상에서 이득이 급감하고, 블록 세분화는 S=2,4,8(N=15,8,4)이 각각 1.845, 1.847, 1.852로 전체 깊이 접근(1.843)에 근접한다.

개발자 관점에서 이 논문의 실용적 가치는 두 가지다. 첫째, 잔차 경로를 손대지 않고 깊이 스케일링 효율을 끌어올리는 모듈이라는 점에서 기존 학습 파이프라인에 비교적 적은 수정으로 얹을 수 있고, 저장·통신이 AttnRes와 동일한 O(Nd)로 유지되므로 파이프라인 병렬화 스케줄을 다시 설계할 필요가 없다. 둘째, m과 블록 수가 성능-비용을 좌우하는 실질적 손잡이라는 점이다. m=4, N≈8이 기본값으로 제시되며 그 이상은 이득이 작다. 다만 이득의 상당 부분이 합성곱 기반 소스 확장에서 나오므로, 시퀀스 축 국소 문맥을 활용할 수 없는 구조(예: 양방향이 아닌 극단적 스트리밍, 매우 짧은 시퀀스)에서는 이점이 줄어들 수 있다는 점을 확인해야 한다.

한계와 전제를 명시한 별도 절은 없다. 저자들이 결론에서 밝히는 바는 다중 상태 표현과 층 간 접근의 이득이 완전히 겹치지 않는다는 관찰, 그리고 더 다양한 모델 구조와 대규모 사전학습 설정에서 두 강점을 결합하는 우아한 잔차 설계가 필요하다는 전망이다. 실험 범위도 제한적이다. 절제 실험은 10B MoE 15층 단일 구성에서만 수행됐고, 스케일링 법칙 적합은 불환원 손실 E=0.72를 고정 추정치로 두는 전제에 의존한다. 블록 세분화에서는 세밀하게 나눌수록 이득이 체감하고, 스트림 수 역시 4를 넘으면 개선이 미미해 기본값 선택이 성능-비용 타협의 결과임을 인정한다. 평가는 DeepSeekMoE 스타일 모델과 영어·중국어 벤치마크에 한정되어 있다.