기술 요약: 정규화인가 국소화인가: 양자화 시 학습 단계의 K/V 캐시 기하학적 구조가 이득을 주는 경우

문제 정의

크로스 엔트로피(CE)로 학습된 트랜스포머 언어 모델은 일반적으로 표현형이 몇 개의 공유된 방향으로 집중되는 비등방성(anisotropic) 은닉 상태 기하학을 발달시킨다. JEPA(Joint-Embedding Predictive Architectures)는 임베딩을 등방성($N(0, I)$)으로 밀어내기 위해 SIGReg(스케치된 Epps–Pulley 정규성 검정)와 같은 목적 함수를 도입했으나, LLM에 SIGReg를 적용하려는 기존의 시도들은 무효한 결과만을 냈다. 이러한 실패는 SIGReg가 비표준 스트림에 보조적으로 적용되었거나, 파인튜닝 과정에서 최종 레이어의 상태에만 적용되었을 때 발생했다.

중요한 공백이 남아 있다: SIGReg가 표준 자기회귀(AR) 사전 학습의 기하학적 구조를 재형성하는가? 더 중요한 것은, 이러한 재형성이 사후 양자화(post-hoc quantization) 과정에서 Key/Value (K/V) 캐시에 이득을 주는가 하는 점이다. 기존 연구는 사후 압축이나 구조적 변경에 집중해 왔으며, 학습 단계의 분포 정규화를 통해 표준 K/V 캐시의 기하학적 구조를 개선하여 양자화 강건성을 높일 수 있는지에 대한 질문은 여전히 미개척 영역으로 남아 있다.

방법론

저자들은 10B FineWeb 토큰으로 사전 학습된 110M 파라미터 규모의 Llama 스타일 모델을 사용하여 3단계 실험 설계를 수행한다.

1단계: 용량-반응 및 은닉 상태 특성 분석

• 목적: 사전 학습 중 CE 손실에 SIGReg를 추가하는 것이 은닉 상태의 비등방성을 줄이는지 결정한다.
• 설정: 저자들은 정규화 가중치 $\lambda$ ($0, 0.001, 0.01, 0.1, 1.0$)를 스윕(sweep)한다.
• 지표: 쌍별 코사인 비등방성(pairwise-cosine anisotropy) 및 퍼플렉시티(perplexity).
• 결과: $\lambda=0.01$이 최소한의 퍼플렉시티 비용으로 비등방성을 유의미하게 감소시키는 최적의 운용 지점으로 식별되었다.

2단계: 브릿지 프로브 (적응 vs. 개조)

기하학적 변화가 K/V 캐시로 어떻게 전파되는지, 그리고 이를 어떻게 달성할 수 있는지 확인하기 위해, $\lambda=0.01$ 체크포인트로부터 시작하는 세 가지 특정 프로브를 테스트한다:

• A1 (적응 - Adaptation): 트렁크(trunk)의 가중치를 해제하고, 모든 레이어에 걸쳐 은닉 상태($h$)와 K/V 투영($K_\ell, V_\ell$) 모두에 SIGReg를 직접 적용하여 학습을 지속한다.
• A2 (고정된 예측 - Frozen Prediction): 고정된 트렁크 상에서 레이어별 예측기를 학습시켜, K/V 기하학이 은닉 상태만으로 복구 가능한지 확인한다.
• A3 (고정된 디코딩 - Frozen Decoding): 고정된 트렁크의 최종 은격 상태로부터 K/V를 재구성하도록 디코더를 학습시킨다.
• 대조군 (Control): 트렁크를 동일한 토큰 예산만큼 지속하되, K/V 정규화 항 없이($\lambda_{KV}=0$) 진행하는 "매치된 지속(matched continuation)" 과정을 거친다.

3단계: 양자화 평가

저자들은 사후 가짜 양자화(양자화 $\to$ 역양자화)를 적용하여 생성된 모델의 K/V 캐시를 평가한다.

• 테스트된 방식: 대칭형 그룹 프리(group-free) 양자화(토큰별 vs. 채널별), 비대칭 그리드(제로 포인트 포함), 그리고 혼합 배치(KIVI 스타일: 채널별 K, 토큰별 V).
• 지표: 모델 자체의 풀 프리시전(full-precision) 참조 모델 대비 $\Delta$NLL (음의 로그 가능도).

주요 결과

1. 은닉 상태 기하학

$\lambda=0.01$에서 SIGReg는 세 개의 쌍을 가진 시드 전체에서 은닉 상태의 쌍별 코사인 비등방성을 38% 감소시킨다. 퍼플렉시티는 0.35% 미만으로 증가하며, 제로샷 능력(LAMBADA, HellaSwag 등)은 일관된 성능 저하 없이 깨끗하게 유지된다.

2. K/V 캐시 기하학 및 "개조(Retrofit)"의 실패

• 자동 전파 불가: 은닉 상태의 비등방성 감소가 K/V 캐시로 자동으로 전파되지 않는다. 고정된 트렁크 프로브(A2 및 A3)는 기하학을 복구하는 데 실패했으며, 이는 K/V 비등방성이 단순히 최종 은닉 상태의 함수가 아님을 나타낸다.
• 직접적인 정규화 필요: 오직 A1(K/V 정규화를 직접 적용한 지속 학습)만이 캐시 비등방성을 줄이는 데 성공하여, 평균 캐시 비등방성을 94% 감소시켰다.
• 메커니즘 격리: K/V 항이 없는($\lambda_{KV}=0$) 매치된 지속 모델은 이 효과를 재현하지 못했으며, 이를 통해 K/V 항이 필수적인 원인임을 격리하여 입증했다.

3. 양자화 이득

정규화된 기하학의 이점은 양자화 방식에 따라 크게 달라진다:

• 거친 스케일 (그룹 프리): 변환되지 않은 대칭형 그룹 프리 양자화 하에서, A1 적응 모델만이 토큰별 스케일링보다 채널별 스케일링을 선호하는 유일한 조건이었다. 3비트 대칭 채널별 양자화에서 베이스라인은 A1 모델보다 4.3~7.9배 더 많은 $\Delta$NLL 손해를 입는다.
• 미세한 스케일 (KIVI 스타일): 전체 시뮬레이션된 KIVI 스타일 구성(혼합 배치, 제로 포인트, 그룹 스케일 $g=32$ 포함) 하에서는 이 이점이 사라진다. 베이스라인을 포함한 모든 모델이 거의 동일한 성능에 도달한다.
• 저장 오버헤드: 그룹 제로 포인트의 메타데이터 오버헤드를 고려하여 3비트 그룹형 저장이 4비트 그룹 프리와 유사한 수준이 되더라도, 토큰 로컬 그룹화 및 혼합 스케일링 구성 하에서는 학습 단계의 이점이 사라진다.

의의 및 주장

본 논문은 표준 K/V 캐시 기하학의 학습 단계 분포 정규화가 사후 캐시 양자화에 미치는 영향을 평가한 최초의 연구라고 주장한다.

• 기하학적 제어: 저자들은 K/V 캐시 기하학이 직접적인 정규화를 통해 학습 중에 명시적으로 제어될 수 있음을 입증했다. 그러나 이 효과는 은닉 상태 정규화만으로는 발생하지 않으며, 고정된 트렁크 프로브를 통해 "개조(retrofit)"될 수도 없다.
• 조건부 유용성: 학습 단계의 개입은 양자화 스케일이 거칠 때(긴 토큰 범위를 포괄할 때)만 실질적인 이득을 제공한다. 이 이점은 정규화된 기하학과 양자화 방식 사이의 상호작용에 특이적이며, 현대적인 미세 조정 구성(제로 포인트 및 로컬 그룹화가 있는 KIVI 스타일 등)에는 일반화되지 않는다.
• 지표 해석: 본 논문은 쌍별 코사인 비등방성 지표가 공분산 화이트닝(covariance whitening)보다는 평균 방향 정렬(공유된 평균에 의해 지배됨)을 측정함을 명확히 한다. 결과는 이 둘의 분리를 보여준다: 즉, 평균 정렬은 감소하지만 스펙트럼 특성(예: top-1 share)은 K와 V 레이어에서 다르게 행동하며, 이는 정규화가 주로 평균/로그 방향(rogue directions)에 작용함을 시사한다.

한계점

저자들은 다음과 같이 주장의 범위를 명시적으로 제한한다:

• 영역: 결과는 110M 파라미터와 FineWeb 코퍼스에 국한되며, 다른 캐시 레이아웃(GQA/MQA)을 가진 더 큰 모델에 대해서는 테스트되지 않았다.
• 시뮬레이션: 양자화는 실제 시간 또는 메모리 측정 없이 시뮬레이션(가짜 양자화)되었다.
• 프로브 범위: 고정된 프로브에 대한 부정적 결과는 특정 용량과 예산에 기반하며, "중간 단계"(예: LoRA 규모의 적응)는 테스트되지 않았다.
• "최초의 등방성" 주장 아님: 저자들은 등방성을 위해 정규화하는 최초의 연구라고 주장하지 않는다(CosReg, SimCTG 등을 인용). 대신, K/V 캐시 기하학을 양자화 강건성을 위해 구체적으로 타겟팅한 최초의 연구임을 강조한다.