Do Value Vectors in Deep Layers Need Context from the Residual Stream?

기술 요약: 가치 뱅크 (Bank of Values, BoV)

문제 정의
현대 대규모 언어 모델(LLM)의 중추인 트랜스포머 아키텍처는 쿼리(query), 키(key), 가치(value) 벡터가 잔차 스트림(residual stream)으로부터 계산되는 어텐션 레이어에 의존합니다. 잔차 스트림은 쿼리와 키를 위해 문맥을 효과적으로 축적하지만, 심층 레이어에서 **가치 벡터(value vectors)**에 문맥이 필요한지에 대한 필요성은 충분히 검토되지 않았습니다. 표준 어텐션은 잔차 스트림($x_i$)을 사용하여 가치 벡터($V = x_i W_V$)를 계산하는데, 이는 원래의 토큰 정보와 축적된 문맥을 혼합합니다. 가치 벡터에 문맥이 없는 토큰 정보를 주입하려는 최근의 연구들(예: Value Residual Learning, SVFormer)은 혼재된 결과를 보여왔으며, 종종 문맥이 없는 성분의 기여를 분리해내지 못하거나 모든 레이어에 비효리적으로 적용하는 문제를 보였습니다. 본 논문이 다루는 핵심 문제는 심층 레이어가 진정으로 문맥 의존적인 가치 벡터를 필요로 하는지, 아니면 잔차 스트림을 투영하는 계산 오버헤드 없이 원래의 문맥이 없는 토큰 정보를 보존하는 것이 더 이득인지 여부입니다.

방법론
저자들은 가치 벡터의 구성을 체계적으로 조사한 후, **가치 뱅크(Bank of Values, BoV)**라고 불리는 새로운 아키텍처 설계를 제안합니다.

1. 체계적 절제 연구 (Systematic Ablation Study):
   저자들은 12개 레이어, 1억 3,500만 파라미터 모델을 사용하여 가치 벡터의 최적 구성을 결정하기 위해 네 가지 축에 걸쳐 실험을 수행했습니다:

   • 가산적(Additive) vs. 대체적(Substitutive): 표준 가치 벡터 전체를 문맥이 없는 소스로 교체하는 것과 이를 선형적으로 결합하는 것을 비교했습니다.
   • 계수(Coefficients): 학습 가능한(unbounded) 계수와 고정된 계수를 테스트했습니다.
   • 문맥 없는 정보의 출처: 초기 임베딩의 레이어별 투영($x_0 W_V$)과 첫 번째 레이어의 공유 가치 벡터($V_1$)를 비교했습니다.
   • 대상 레이어: 이러한 변화를 모든 레이어, 매 두 번째 레이어, 또는 마지막 1/3의 레이어에 적용했을 때의 영향을 평가했습니다.

   절제 연구의 주요 결과:

   • 심층 레이어는 문맥 의존적인 가치 벡터($V$)보다 문맥이 없는 성분($\tilde{V}$)을 강력하게 선호합니다. 가산적 변형에서 $\tilde{V}$에 대한 학습된 계수는 $V$에 대한 계수(종종 0으로 수렴함)보다 훨씬 컸습니다(종종 >10).
   • $V$를 $\tilde{V}$로 대체하는 것은 마지막 1/3의 레이어에서만 유의미한 성능 향상을 가져왔으며, 얕은 레이어에 적용하면 성능이 저하되었습니다.
   • 일단 심층 레이어가 고품질의 문맥 없는 성분에 접근하게 되면, 문맥 의존적 성분을 추가하는 것은 미미한 이점만을 제공했습니다.
   • 레이어별 문맥 없는 벡터($x_0 W_V$)는 공유 벡터($V_1$)보다 우수한 성능을 보였는데, 이는 각 레이어가 토큰을 고유한 부분 공간으로 투영할 수 있게 해주기 때문입니다.

2. 가치 뱅크 (Bank of Values, BoV) 아키텍처:
   심층 레이어가 결정론적이고 입력 시퀀스와 독립적인 문맥 없는 가치 벡터로부터 이득을 얻는다는 발견에 기반하여, 저자들은 BoV를 제안합니다.

   • 메커니즘: 마지막 1/3의 레이어에서 표준 밀집 가치 행렬 투영($x_i W_V$)은 희소 룩업 테이블(sparse lookup table)($E_V$)로 대체됩니다.
   • 구현: 토큰 ID가 $i$인 경우, 가치 벡터는 $v_i = \gamma_v E_V[i]$로 직접 검색됩니다. 여기서 $E_V \in \mathbb{R}^{|V| \times d}$는 모든 어휘 사전의 토큰에 대해 사전 계산된 가치 벡터를 저장하는 학습 가능한 파라미터 테이블이며, $\gamma_v$는 학습 가능한 스케일링 계수입니다.
   • 추론 효율성: 모든 이전 토큰에 대해 가치 벡터를 캐싱해야 하는 표준 어텐션(레이어당 $O(Td)$ 메모리 필요)과 달리, BoV는 값을 캐싱하지 않습니다. 대신, 순방향 패스 중에 정적 파라미터 테이블 $E_V$로부터 $O(Td)$ gather 연산을 수행합니다. 테이블 $E_V$는 호스트 메모리로 오프로드될 수 있으며, 필요한 항목만 프리페치(prefetch)됩니다.

주요 기여

• 통찰: 본 논문은 심층 트랜스포머 레이어의 가치 벡터가 주로 문맥이 없는 원래의 토큰 정보로부터 이득을 얻으며, 잔차 스트림에서 유도된 문맥 의존적 성분은 대체로 중복된다는 점을 입증합니다.
• 아키텍처: 토큰별 테이블로부터의 정적 룩업으로 심층 레이어의 가치 벡터 계산을 대체하는 방법인 **가치 뱅크 (BoV)**를 도입합니다.
• 효율성: BoV는 심층 레이어를 위한 가치 벡터를 계산하고 캐싱할 필요를 없애, 긴 문맥 추론 시 메모리 사용량을 줄이고 토큰당 FLOPs를 낮춥니다.

결과
저자들은 계산 제어 예산($1.50 \times 10^{18}$ 및 $3.91 \times 10^{19}$ FLOPs) 하에 훈련된 1억 3,500만 및 7억 8,000만 파라미터 규모의 모델에 대해 BoV를 평가했습니다.

• 검증 손실 (Validation Loss): BoV는 두 모델 규모 모두에서 홀드아웃 분할된 ClimbMix 데이터셋에 대해 표준 어텐션 베이스라인보다 낮은 검증 손실을 달s성했습니다.
• 벤치마크 성능: 7억 8,000만 규모에서, BoV는 표준 어텐션 베이스라인과 비교하여 21개 DCLM CORE 벤치마크의 평균 점수를 개선했습니다.
• 비교: BoV는 토큰 정보를 가치 벡터에 주입하는 기존 방법들(예: Value Residual Learning)과 대등하거나 이를 능가하면서도 더 적은 메모리와 더 적은 FLOPs를 사용했습니다. 구체적으로, 문맥 의서적 가치를 계산하고 캐싱하는 비용을 여전히 부담하는 방식보다 우수한 성능을 보였습니다.

의의
본 논문은 BoV가 심층 레이어에서 가치 벡터를 계산하는 데 문맥이 필요하지 않다는 통찰을 활용함으로써, 표준 어텐션에 대한 더 효율적인 대안을 제공한다고 주장합니다. 가치 벡터를 동적인 계산이 아닌 희소 모델 파라미터로 유지함으로써, BoV는 모델 성능을 희생하지 않으면서도(때로는 향상시키면서도) 긴 문맥 추론을 위한 메모리 발자국을 줄이고 계산 비용을 낮춥니다. 이는 모든 어텐션 벡터를 잔차 스트림으로부터 계산한다는 기본 패러다임에 도전하며, 더 효율적인 트랜스포머 아키텍처를 향한 경로를 제시합니다.