Stuck on "A": Diagnosing and Repairing Interface Injury in Attention-to-KDA Linearization of a 0.6B Language Model

기술 요약: Attention-to-KDA 선형화에서의 인터페이스 손상 진단 및 수리

문제 정의
본 논문은 사전 학습된 풀 어텐션(full-attention) 언어 모델을 선형 어텐션 아키텍처, 특히 Kimi Delta Attention(KDA) 메커니즘으로 변환할 때 발생하는 "손실적(lossy)" 특성을 조사한다. 선형 어텐션은 긴 컨텍스트 추론 및 에지 배포를 위한 구조적 이점(O(n) KV 캐시 대신 O(1) 순환 상태)을 제공하지만, 변환 과정은 표준 지표가 성공을 시사함에도 불구하고 종종 작업 성능을 보존하는 데 실패한다.

저자는 Qwen3-0.6B-Base(28개 레이어)를 21개 레이어는 KDA로 교체하고 7개 레이어는 Grouped Query Attention(GQA)을 유지하는 하이브리드 아키텍처로 변환하는 데 집중한다. 이는 단일 32GB GPU와 수십억 개가 아닌 수천만 개의 토큰이라는 엄격한 소비자급 제약 조건 하에서 수행된다.

핵심적인 진단 질문은 다음과 같다: 변환 과정에서 정확히 무엇이 망가지는가? 표준 평가는 역설을 보여준다: 레이어별 정렬(layer-wise alignment)과 엔드 투 엔드 KL 증류(end-to-end KL distillation)를 거친 후, 학생 모델은 교사 모델에 근접한 퍼플렉시티(Cross-Entropy 차이 +0.128 nats)를 달성했지만, 다지선다형(multiple-choice) 정확도에서는 무작위 확률 수준(~25–29% on C-Eval, 교사의 50.6% 대비)에 머물며 붕괴했다.

방법론
저자는 이 성능 격차를 진단하고 수리하기 위해 다단계 파이프라인을 제안한다:

1. 하이브리드 아키텍처 및 수술 (Hybrid Architecture & Surgery):

   • 레이아웃: 3, 7, 11, 15, 19, 23, 27번 레이어는 네이티브 GQA를 유지하며, 나머지 21개 레이어는 KDA(16개 헤드, 헤드 차원 128)로 변환된다.
   • KDA 메커니즘: 망각 게이트($\alpha_t$)와 쓰기 게이트($\beta_t$)를 갖춘 퍼 채널 게이트 델타 규칙(per-channel gated delta rule)을 구현한다.
   • 초기화: 결정적인 절제 연구(ablation study)를 통해 단순한 프로젝션 이식(projection transplant)은 실패함을 확인했다. 성공적인 구성(g6)은 순환 상태를 안정화하기 위해 특정 감쇠 스케일(decay scales)과 바이어스 설정을 사용한다.

2. 학습 단계:

   • 단계 2 (레이어별 정렬): 고정된 교사의 은닉 상태(hidden states)와 KDA 은닉 상태를 MSE를 사용하여 정렬한다. 이는 Cross-Entropy를 9.48에서 4.13으로 줄였으나, C-Eval 정확도는 25.5%에 머물렀다.
   • 단계 3a (엔드 투 엔드 KL 증류): 교사와 학생 분포 사이의 forward KL divergence를 사용하여 사전 학습 텍스트로 학습한다. 이는 CE 차이를 +0.128 nats로 좁혔으나, C-Eval 정확도는 28.8%로만 상승했다.
   • 단계 3b (포맷 타겟형 수리): 교사가 검증한 다지선다형 문제(MCQs), 시(poetry), 번역 QA를 사용하는 1,000-스텝 "완성 전용(completion-only)" KL 단계를 수행한다. 결정적으로, 이 단계는 일반적인 텍스트 생성이 아닌 답변의 포맷에 집중했다.

3. 4-순열 진단 (The Four-Permutation Diagnostic):
   저자는 퍼플렉시티를 넘어 진단하기 위해 161개의 깨끗한 MCQ(벤치마크와 중복 없음)를 사용하는 진단 프로토콜을 설계했다. 각 질문은 답변 옵션(A, B, C, D)의 네 가지 순환 회전(cyclic rotations) 하에서 평가되었다.

   • 지표: "스티키니스(Stickiness)"는 모델이 콘텐츠의 위치와 상관없이 동일한 라벨(예: "A")을 얼마나 자주 예측하는지를 측정한다.
   • 결과: 단계 3a 모델은 심각한 **인터페이스 손상(interface injury)**을 보였다. 모델은 "A"를 81%의 확률로 예측했으며, 161개 중 106개 질문에 대해 네 가지 회전 모두에서 동일한 라벨을 유지했다. 즉, 모델은 지식을 옵션 라벨에 매핑하는 능력을 상실하여, 선택지의 내용을 무시하고 특정 라벨을 기본값으로 출력하게 되었다.

4. 페르소나 정렬 (Persona Alignment):
   수리 후, 모델은 주입된 페르소나("Qingyi")와 함께 SFT를 거쳤으며, 변환이 모델의 정체성이나 대화 능력을 저하시키지 않도록 한 차례의 온-폴리시(on-policy) DPO를 수행했다.

주요 결과

• "퍼플렉시티의 거짓말 (The Perplexity Lie)": 본 연구는 은닉 상태 정렬과 낮은 KL divergence가 작업 능력을 보장하지 않음을 확인한다. 모델은 교사의 분포 통계는 모방할 수 있지만, 그 지식을 특정 작업을 위해 활용하는 데는 실패할 수 있다.
• 인터페이스 손상의 진단: 4-순열 테스트는 모델이 균일한 지식 손실(무작위 추측)을 겪는 것이 아니라, **라벨 스티키니스(label-stickiness)**를 겪고 있음을 드러냈다. 모델은 답변 콘텐츠를 추적하지 못하고 특정 옵션 라벨을 기본값으로 선택했다.
• 수리 효능: 포맷 타겟형 단계 3b 수리는 다음과 같은 결과를 가져왔다:
  • C-Eval 정확도를 28.83%에서 41.31%로 향상시켰다.
  • "A" 예측 점유율을 81.06%에서 58.5%로 감소시켰다.
  • "4개 회전 간 동일 라벨" 지표를 절반으로 줄였다 (106/161 $\to$ 47/161).
  • 정답 옵션 마진(correct-option margin)을 음수(−0.117 nats)에서 양수(+0.233 nats)로 전환했다.
  • 트레이드오프: 이 수리는 일반적인 분포 적합도(Valid CE/KL)를 약간 증가시키는 비용을 치렀으며, 이는 모델이 MCQ 포맷에 특화되었음을 시사한다.
• 최종 성능: 페르소나 정렬 후, 최종 모델은 C-Eval에서 **41.83%**를 달성하였으며, 이는 pre-SFT 수리 단계의 노이즈 밴드 내에 머물러 있어, 변환 과정에서 치명적인 망각이 발생하지 않았음을 나타낸다.

엔지니어링 교훈 및 실패 사례
본 논문은 과정 중 겪은 몇 가지 중요한 엔지니어링 함정을 기록한다:

• FP32 마스터 실패 모드: 파라미터와 Adam 옵티마이저 모멘트에 모두 bf16을 사용했을 때, 작은 가중치 업데이트가 "조용히 삼켜지는"(sub-ULP updates) 현상이 발생하여, 그래디언트 노름은 정상임에도 가중치가 고정되는 현상이 발생했다. 해결책은 FP32 마스터 가중치를 사용하는 것이었다.
• 캐시 패리티 (Cache Parity): 새로운 KDA 레이어가 초기에는 증분 캐시 업데이트를 무시하여, 두 번째 토큰부터 생성이 쓰레기 값으로 붕괴되는 문제가 있었다. 특정 캐시 구현 수정으로 재학습 없이 이를 해결했다.
• 잘못된 이식 (False Transplants): 교사의 Q/K/V/O 프로젝션을 이식하려던 초기 시도들은 기능적으로 실패했다. 성공적인 접근법은 프로젝션 이식을 완전히 포기하는 것이었다.

의의 및 주장
본 논문은 네 가지 주요 기여를 주장한다:

1. 진단 도구: "인터페이스 손상"(지식을 출력 포맷에 매핑하는 능력 상실)과 "지식 손실"(균일한 망각)을 구분하는 4-순열 프로토콜을 제공하며, 이는 모든 선형화 또는 압축 파이프라인에 적용 가능하다.
2. 재현 가능한 레시피: 소비자 예산 내에서 Qwen3를 KDA 하이브리로 변환하는 완전한 레시피를 제공하며, 공개 기록에서 흔히 누락되는 특정 실패 모드(게이트 절제, 캐시 버그, 업데이트 삼킴 현상)를 기록한다.
3. 페르소나 생존: 페르소나 정렬(SFT + DPO)이 대규모 재학습이나 치명적 망각 없이 아키텍처 변환 과정에서 "거의 공짜로" 생존할 수 있음을 입증한다.
4. 오픈 릴리스: 코드, 가중치, 레시피 및 실패 기록을 포함한 전체 감사 추적(audit trail)을 공개한다.

저자는 소비자 예산으로 0.6B 모델을 KDA 하이브리드로 변환하는 것이 가능하지만, 표준 메트릭 스택(퍼플렉시티, KL, 은닉 상태 MSE)이 손상을 체계적으로 잘못 보고한다는 결론을 내린다. 성공의 주요 장벽은 반드시 내부 지식의 손실이 아니라, 그 지식을 특정 출력 포맷(예: 다지선다형 라벨)으로 매핑하는 인터페이스의 저하이며, 이는 타겟팅된 포맷 감독을 통해 수리될 수 있다.