기술 요약: 역할 분리형 어텐션 잔차 (Role-Decoupled Attention Residuals)

문제 정의

표준 트랜스포머 아키텍처에서 잔차 스트림(residual stream)은 각 레이어가 직전의 상태를 직렬 합산 방식으로 상속받는 고정된 메모리 역할을 합니다. 최근의 깊이 라우팅(depth-routing) 아키텍처인 **블록 어텐션 잔차(Block Attention Residuals, Block AttnRes)**는 하위 레이어가 단순히 바로 이전의 상태가 아니라, 이전 표현들의 콘텐츠 의존적 혼합물(content-dependent mixture)을 검색할 수 있게 함으로써 이를 개선합니다. 이는 깊이를 고정된 계산 순서에서 학습된 검색 축(learned retrieval axis)으로 변환합니다.

그러나 Block AttnRes는 특정 설계 제약을 도입합니다. 즉, 세 가지 어텐션 투영(Query, Key, Value) 모두에 대해 공유된 입력을 구성하기 위해 **단일 깊이 분포(single depth distribution)**를 사용한다는 점입니다. 이는 다음과 같이 기능적으로 구별되는 두 역할을 결합합니다:

1. 매칭 기하학(Matching Geometry): Query와 Key에 의해 결정되며, 어디를 볼 것인지(어떤 토큰 위치가 일치하는지)를 결정합니다.
2. 콘텐츠 검색(Content Retrieval): Value에 의해 결정되며, 무엇을 읽을 것인지(어떤 콘텐츠를 가져올지)를 결정합니다.

본 논문은 "어디를 볼지"를 결정하는 데 최적인 표현이 "무엇을 읽을지"를 결정하는 데도 반드시 최적이어야 한다는 가정이 불필요한 결합임을 제기합니다. 핵심 질문은 다른 모든 요소가 동일할 때, 이 두 역할을 분리하는 것이 언어 모델 사전 학습(pretraining)을 개선하는가 하는 점입니다.

방법론: 역할 분리형 어텐션 잔차 (RD-AttnRes)

저자는 Block AttnRes의 최소한의 아키텍처 완화 모델인 **역할 분리형 어텐션 잔차(RD-AttnRes)**를 소개합니다.

• 아키텍처: RD-AttnRes는 부모 아키텍처의 소스 세트, 블록 토폴로지, 정규화, MLP 경로 및 쓰기백(writeback) 메커니즘을 그대로 유지합니다. 유일한 수정 사항은 어텐션 측의 깊이 읽기(depth read) 과정에서 발생합니다.
• 분리 메커니즘:
  • Queries와 Keys (Q/K): 단일 학습된 경로 쿼리($q^{QK}_\ell$)를 공유하여 잔차 깊이 메모리로부터 혼합물($h^{QK}_\ell$)을 계산합니다. 이는 일관된 매칭 기하학을 보존합니다.
  • Values (V): 독립적인 경로 쿼리($q^V_\ell$)를 학습하여 동일한 잔차 소스 세트로부터 별도의 혼합물($h^V_\ell$)을 계산합니다.
• 파라미터 효율성: 이 수정 사항은 레이어당 차원 $d$(모델 너비)를 갖는 단 하나의 벡터만을 추가합니다. 테스트된 규모(120M 및 343M 파라미터)에서 이는 파라미터 증가량이 미미히(<0.01%) 함을 의미합니다.
• 정확한 복구(Exact Recovery): 만약 Value 경로 쿼리가 Q/K 쿼리와 동일하게 설정된다면($q^V_\ell = q^{QK}_\ell$), RD-AttnRes는 정확히 부모 모델인 Block AttnRes로 환원됩니다. 이는 본 제안이 여러 변화를 동시에 일으키는 시스템이 아니라, 엄격하고 검증 가능한 확장임을 보장합니다.
• 계산 비용: 이 방법은 어텐션 레이어당 $O(BTK_\ell d)$의 연산량(여기서 $B$는 배치 크기, $T$는 시퀀스 길이, $K_\ell$은 소스의 수)을 추가하지만, 새로운 $O(T^2d)$ 토큰 간 어텐션 항을 도입하지는 않습니다.

실험 설계

본 연구는 역할 분리를 격리하기 위해 엄격한 쌍을 이룬 실험 프로토콜을 채택합니다:

• 데이터셋: FineWeb-Edu (10BT 샘플), 학습, 검증, 테스트를 위한 결정론적 분할을 사용합니다.
• 모델: 두 가지 디코더 전용 Pre-Norm 트랜스포머 규모:
  • S: 120M 파라미터 (16 레이어, 너비 576).
  • M: 343M 파라미터 (24 레이어, 너비 832).
• 프로토콜: 각 규모당 5개의 매칭된 시드(seed)를 사용합니다. Block AttnRes와 RD-AttnRes는 동일한 초기화, 데이터 순서, 최적화 스케줄, 그리고 2.0B 토큰 예산을 가지고 학습되었습니다.
• 대조군(Controls): 역할 분리의 효과를 격리하기 위한 방향성 대조군들을 0.5B 토큰 시점에서 설정했습니다:
  • Tied (결합됨): Value 경로가 Q/K 경로와 일치하도록 강제함 (부모 모델을 복구).
  • Average (평균): 두 개의 경로를 실행하되 그 혼합물을 평균함 (이중 실행 자체의 효과를 테스트).
  • Frozen-V (고정된 V): Value 경로를 균일하게 유지함 (고정된 경로가 도움이 되는지 테스트).
• 진단: Q/K와 V의 깊이 분포 간의 젠슨-샤논(Jensen–Shannon, JS) 발산(divergence)을 측정하여 경로가 서로 다른 전략을 학습하는지 확인합니다.

주요 결과

1. 검증 손실(Validation Loss) 개선: RD-AttnRes는 모든 10개의 쌍 비교(5개 시드 × 2개 규모)에서 검증 음의 로그 가능도(NLL)를 개선했습니다.
   • 120M 규모: 평균 NLL 감소량은 0.0301이며, 이는 퍼플렉서티(perplexity)의 2.97% 감소에 해당합니다.
   • 343M 규모: 평균 NLL 감소량은 0.0247이며, 이는 퍼플렉서티의 2.43% 감소에 해당합니다.
2. 이득의 지속성: 개선 효과는 모든 평가 마일스톤(0.25B, 0.5B, 1.0B, 2.0B 토큰)에서 관찰되지만, 학습이 진행됨에 따라 격차의 절대적인 크기는 약간 좁아집니다.
3. 대조군 검증:
   • Tied와 Average 대조군은 부모 모델인 Block AttnRes와 유사한 성능을 보였으며, 이는 이득이 파라미터 수나 중복 계산 때문이 아님을 나타냅니다.
   • Frozen-V는 성능이 저조했으며, 이는 Value 경로가 고정된 것이 아니라 반드시 학습되어야 함을 확인시켜 줍니다.
4. 경로 발산: 학습된 Q/K 및 V 분포는 동일한 솔루션으로 붕괴되지 않습니다. 레이어 전체의 평균 JS 발산은 120M에서 약 0.218, 343M에서 약 0.215로 나타났으며, 이는 매칭과 콘텐츠 검색을 위해 서로 다른 깊이 읽기가 학습되었음을 확인해 줍니다.

의의 및 주장

본 논문은 보편적인 주장이 아닌 좁고 제한된 범위의 주장을 합니다. 저자는 RD-AttnRes가 모든 최근의 잔차 라우팅 아키텍처보다 우월하다거나 보편적인 스케일링 동작을 확립한다고 주장하지 않습니다.

대신, 결과는 테스트된 학습 환경(FineWeb-Edu, 2B 토큰, 특정 모델 규모) 내에서, Block AttnRes의 공유 QKV 깊이 읽기가 중요한 제약 사항임을 뒷받립합니다. 증거는 매칭(Q/K)과 콘텐츠 검색(V)이 서로 다른 깊이 표현을 필요로 한다는 가설을 지지합니다.

인정된 한계점:

• 본 연구는 하나의 영어 웹 코퍼스와 10억 미만 파라미터 모델에 국한되어 있습니다.
• 절대적인 성능 격차는 더 긴 학습 기간 동안 좁아지며, 향-후 연구를 통해 4B~10B 토큰에서도 이 격차가 유지되거나 커지는지 확인할 필요가 있습니다.
• 시스템 측면의 주의사항: 현재 구현은 참조 모델인 Block AttnRes에 비해 처리량(throughput) 저하(~9–10%)와 메모리 사용량 증가를 초래하지만, 퓨즈드 커널(fused kernels)을 통해 완화될 수 있습니다.
• 다른 최근 아키텍처(예: MUDDFormer, Delta AttnRes)와의 비교는 개념적인 수준이며, 일치된 경험적 베이스라인은 아닙니다.

결론:
본 논문은 깊이가 학습된 검색 축이 될 때, 공유(sharing)는 구현의 편의성이 아니라 다운스트림 기능(역할)을 따라야 함을 입증합니다. 최소한의 파라미터 오버헤드로 Value 경로를 분리함으로써, RD-AttnRes는 사전 학습 손실을 일관되게 개선하였으며, 이는 매칭 기하학과 콘텐츠 검색이 서로 다른 깊이 표현을 요구하는 별개의 작업임을 검증합니다.