🤖 AI

Beyond Output Matching: Preserving Internal Geometry in NVFP4 LLM Distillatio

이 논문은 출력 매칭만을 수행하는 양자화 인식 증류(QAD)가 NVFP4 모델의 내부 표현 기하학을 보존하는 데 실패하여 추론 및 코딩 작업에서 성능 저하를 초래한다는 점을 식별하고, 내부 구조를 유지함으로써 정확도를 회복하기 위해 레이어별 그람 행렬(Gram matrices)을 정렬하는 방법인 CKA-QAD를 제안한다.

원저자: Fangbo Tu, Junhua Zhao, Chi Liu, Xin Chen, Haifeng Wu, Jian Wan, Srinivasan Manoharan

게시일 2026-06-05

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Fangbo Tu, Junhua Zhao, Chi Liu, Xin Chen, Haifeng Wu, Jian Wan, Srinivasan Manoharan

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신에게 완벽하고 복잡한 요리를 만들 수 있는 천재적인 마스터 셰프(스승)가 있다고 상상해 보세요. 이제 당신은 학생 셰프(학생)에게 똑같은 요리를 가르치고 싶지만, 한 가지 큰 제약이 있습니다. 학생은 아주 작고 품질이 낮은 계량컵과 스푼만을 사용할 수 있다는 것입니다(이것은 모델의 크기를 줄여 더 빠르고 저렴하게 사용하기 위한 방법인 NVFP4 양자화입니다).

보통 학생을 가르칠 때, 당신은 그저 이렇게 말할 것입니다. "최종적인 맛이 내 요리와 똑같아야 해." AI의 세계에서는 이를 출력 매칭(Output Matching, 또는 KL Divergence)이라고 부릅니다. 당신은 최종 결과물(다음 단어의 확률인 "로짓")을 확인하며 학생에게 이렇게 말합니다. "네 요리의 맛이 내 것과 같다면, 너는 잘하고 있는 거야."

문제점: "가짜" 일치

이 논문의 저자들은 이 방식에서 놀라운 결함을 발견했습니다.

그들은 학생 셰프가 속임수를 쓸 수 있다는 것을 발견했습니다. 학생은 재료와 조리 단계를 완전히 이상하고 혼란스러운 방식으로 재배열할 수 있지만, 어찌 된 일인지 최종적인 맛은 마스터 셰프의 요리와 일치하게 만들 수 있었습니다.

현실: 학생의 내부 주방(내부 기하학)은 엉망진창이었습니다. 재료는 잘못된 순서로 섞였고, 열은 이상하게 가해졌으며, "풍미 지도"는 망가져 있었습니다.
환상: 최종적인 맛이 맞았기 때문에, 스승은 학생이 배우고 있다고 생각했습니다. 하지만 학생이 새로운 복잡한 요리(예: 수학 문제나 코딩)를 시도했을 때, 그들은 실패했습니다. 왜냐하면 그들의 내부 주방이 무질서했기 때문입니다.

논문에서는 이를 **"출력 매칭이 내부 퇴화를 가리는 현상(Output Matching masking Internal Degradation)"**이라고 부릅니다. 겉보기에는 좋아 보이지만, 그 내부의 뇌는 뒤섞여 버린 상태입니다. 이는 특히 모델이 깊게 "생각하도록(강화 학습)" 훈련되었을 때, 즉 내부 논리가 매우 중요할 때 더욱 심각하게 나타났습니다.

해결책: CKA-QAD (The "Kitchen Blueprint" Check)

이를 해결하기 위해 저자들은 CKA-QAD라는 새로운 교수법을 도입했습니다.

단순히 최종적인 맛을 확인하는 대신, 그들은 다음과 같은 규칙을 추가했습니다. "너의 주방 구조도 나의 것과 같아야 한다."

그들은 **CKA(Centered Kernel Alignment)**라는 수학적 도구를 사용했습니다. CKA를 설계도 스캐너라고 생각하면 됩니다. CKA는 학생이 다른 냄비나 팬을 사용하는지 여부는 신경 쓰지 않지만(미세한 규모의 변화는 무시함), 재료 간의 관계가 동일한지는 확인합니다.

만약 마스터 셰프가 소금을 후추보다 먼저 넣는다면, 학생도 똑같이 해야 합니다.
만약 마스터 셰프가 특정 패턴으로 밀가루와 달걀을 섞는다면, 학생도 그 패턴을 모방해야 합니다.

이 "설계도 체크"를 훈련에 추가함으로써, 학생은 단순히 맛을 흉내 내는 것이 아니라 내부 주방을 체계적으로 유지하도록 강요받습니다.

실험 결과는 어떠했는가?

연구진은 이 방법을 두 가지 서로 다른 AI 모델(하나의 거대한 하이브리드 모델과 하나의 작은 추론 모델)에 테스트했습니다.

기존 방식 (맛만 확인): 학생은 최종 답변을 흉내 낼 수는 있었지만, 내부의 "뇌"는 스승으로부터 멀어져 갔습니다. 어려운 수학이나 코딩 문제를 풀도록 요청받았을 때, 그들은 비틀거렸습니다.
새로운 방식 (맛 + 설계도): 학생은 맛을 맞추는 동시에 내부 주방을 체계적으로 유지하는 법을 배웠습니다.
- 결과: 학생은 단순히 정답을 맞히는 것을 넘어, 수학(AIME25) 및 코딩(LiveCodeBench)과 같은 복잡한 과제에서 더 나은 성과를 냈습니다.
- 효율성: 이 새로운 방식은 훈련 과정에 약 0.5%의 추가 시간과 약 7%의 메모리만을 추가했습니다. 훨씬 더 똑똑한 학생을 얻기 위한 아주 작은 대가였습니다.

핵심 요점

이 논문은 AI 모델을 압축할 때(줄일 때), 단순히 최종 결과만 봐서는 안 된다고 결론짓습니다. 반드시 모델의 내부 구조가 원본과 동일하게 유지되도록 해야 합니다.

간단히 말해서: 만약 당신이 작은 AI를 똑똑하게 만들고 싶다면, 단지 "정답을 맞춰라"라고 말해서는 안 됩니다. "올바른 방식으로 생각하라"고도 말해야 합니다. 모델이 무엇을 말하는지(결과)뿐만 아니라 어떻게 생각하는지(내부 기하학)를 확인함으로써, 우리는 실제로 깊은 추론 능력을 갖춘 작고 빠른 AI 모델을 만들 수 있습니다.

기술 요약: NVFP4 LLM 증류에서의 출력 매칭을 넘어선 접근법

1. 문제 정의

대규모 언어 모델(LLM)이 지연 시간(latency)과 메모리 제약이 있는 프로덕션 환경에 배포됨에 따라, 저정밀도 추론 포맷, 특히 NVFP4(미세한 블록 스케일링을 갖춘 4비트 부동 소수점 포맷)에 대한 수요가 증가하고 있습니다. 사후 학습 양자화(PTQ) 과정에서 발생하는 정확도 손실을 복구하기 위해, 양자화 인식 증류(Quantization-Aware Distillation, QAD)가 KL 발산(KL divergence)을 통해 양자화된 학생 모델의 출력 분포를 전정밀도(full-precision) 교사 모델과 정렬하는 견고한 방법으로 부상했습니다. 그러나 본 논문은 기존 패러다임의 치명적인 결함을 지적합니다.

저자들은 출력 수준의 정렬만으로는 충실한 능력 복구가 불가능하다고 주장합니다. 이들의 분석에 따르면, 오직 KL 발산만을 최적화하는 표준 QAD는 심각한 **내부 표현 퇴화(internal representational degradation)**를 은폐할 수 있습니다. 구체적으로, 학생 모델이 교사 모델의 출력 로짓(logits)을 성공적으로 일치시키더라도, 중간 단계의 활성화 기하 구조(activation geometries)는 의도치 않은 재편(drift)을 겪게 됩니다. 이러한 "지름길 학습(shortcut learning)"은 특히 강화 학습(RL) 사후 학습을 거친 모델에서 더욱 심각하게 나타나며, 이 경우 복잡한 추론 및 코딩 작업을 수행하는 데 필요한 내부의 의미론적 계층 구조가 출력 정확도가 높음에도 불구하고 훼손됩니다.

2. 방법론: CKA-QAD

출력 매칭과 내부 기하 구조 사이의 괴리를 해결하기 위해, 저자들은 중심 커널 정렬(Centered Kernel Alignment, CKA) 정규화 항을 표준 QAD에 추가한 CKA-QAD를 제안합니다.

핵심 메커니즘

로짓 수준의 증류에만 의존하는 대신, CKA-QAD는 교사와 학생의 **레이어별 그람 행렬(Gram matrices)**을 정렬하는 경량 정규화 항을 도입합니다.

CKA 지표: 이 방법은 중심화된 활성화 행렬( $\bar{H}_t$ 및 $\bar{H}_s$ ) 사이의 유사도를 측정하는 선형 CKA를 활용합니다. CKA는 직교 변환(orthogonal transformations)과 등방성 스케일링(isotropic scaling)에 대한 불변성을 가지므로, NVFP4 양자화에 내재된 비정상적(non-stationary) 활성화 분포와 스케일 변화에 강건합니다.
목적 함수: 전체 손실은 표준 Top-K KL 발산 손실( $L_{TopK-KL}$ )과 동적 C로 구성됩니다.
$L_{total} = L_{TopK-KL} + \text{sg}\left(\frac{L_{TopK-KL}}{L_{CKA} + \epsilon}\right) L_{CKA}$
stop-gradient (sg) 연산자의 사용과 동적 스케일링을 통해, 전역 가중치 $\lambda$ 에 대한 수동 하이퍼파라미터 튜닝 없이 CKA 항을 로짓 손실과 균형 있게 조절할 수 있습니다.

구현 세부 사항

레이어 선택: CKA는 매칭된 디코더 블록(post-attention/MLP/Mamba 업데이트)의 출력을 대상으로 계산됩니다. 저자들은 잔차 스트림(residual stream)의 스케일과 자기회귀적 의존성으로부터 발생하는 아티팩트를 방지하기 위해 pre-normalization 활성화, 어텐션 가중치, KV 캐시 상태를 명시적으로 제외했습니다.
엔지니어링 최적화: 많은 레이어에 걸쳐 CKA를 계산할 때 발생하는 계산 오버헤드를 관리하기 위해 다음을 적용했습니다:
1. Top-K 증류: KL 손실은 교사의 상위 8192개 로짓에 대해서만 계산됩니다.
2. 특징 공간(Feature-Space) CKA: CKA는 명시적인 토큰 간 그람 행렬 대신 특징 공분산 행렬( $\bar{H}^\top \bar{H}$ )을 통해 계산되어 메모리 복잡도를 줄입니다.
3. Hook 기반 캡처: 활성화는 포워드 훅(forward hooks)을 통해 캡처되며, 교사의 활성화는 계산 그래프에서 분리(detach)됩니다.

3. 주요 기여

QAD 실패 모드 진단: 본 논문은 KL 전용 증류가 특히 중간 깊이의 Transformer FFN 레이어와 Mamba 상태 블록에서 표현 드리프트(representational drift)를 악화시킬 수 있음을 보여주는 표현 수준의 진단을 제공합니다. 이 드리프트는 추론 및 코딩 작업의 성능 병목 현상과 상관관계가 있습니다.
CKA 가이드 NVFP4 QAD: 저비트 증류 과정에서 내부 표현 기하 구조를 명시적으로 보존하는 방법인 CKA-QAD를 도입했습니다. 이를 통해 최소한의 오버헤드로 BF16 수준에 근접한 정렬을 복구합니다.
실증적 검증: 두 가지 서로 다른 아키텍처에 대해 종합적인 평가를 수행했습니다:
- Nemotron 3 Nano: RL 사후 학습이 적용된 하이브리드 Mamba–Transformer MoE 모델.
- Qwen3-4B-Thinking-2507: 컴팩트한 Transformer 추론 모델.

4. 실험 결과

저자들은 CKA-QAD를 BF16 교사 모델, NVFP4 PTQ 베이스라인, 그리고 표준 KL 전용 QAD와 비교 평가했습니다.

표현 정렬 (Representational Alignment)

Nemotron 3 Nano: 표준 QAD(KL 전용)는 PTQ(0.993)에 비해 평균 CKA(0.958)가 크게 감소했으며, 마지막 레이어에서 심각한 붕괴(0.740)를 보였습니다. CKA-QAD는 평균 CKA를 0.994로, 마지막 레이어 CKA를 0.985로 복구하여 BF16 교사와 거의 일치하는 수준을 보였습니다.
Qwen3-4B: 유사한 경향이 관찰되었습니다. KL 전용 QAD는 평균 CKA를 0.98로, 마지막 레이어 CKA를 0.86으로 떨어뜨렸습니다. CKA-QAD는 이를 각각 0.99와 0.97로 복구했습니다.

다운스트림 정확도

CKA-QAD는 복잡한 추론 및 코딩 벤치마크에서 일관된 개선을 입증했습니다:

AIME25 (수학 추론): Qwen3-4B에서 68.5% (KL 전용)에서 72.3% (CKA-QAD)로 향상되었습니다.
GPQA-D (전문 지식): 59.5%에서 **61.1%**로 향상되었습니다.
LiveCodeBench-v5 (코드 생성): 57.9%에서 **59.8%**로 향상되었습니다.
Nemotron 3 Nano: 유사한 회복을 보였으며, AIME25는 87.9%에서 **88.3%**로, LiveCodeBench는 68.9%에서 **70.3%**로 향상되었습니다.

효율성 오버헤드

이 방법은 완만한 학습 비용을 도입합니다:

실제 실행 시간(Wall-clock time): 스텝당 +0.5%.
피크 VRAM: +7.0%.
처리량(Throughput): -0.5%.
수렴(Convergence): 목표 KL에 도달하는 스텝 수가 약간 감소했습니다 (-3.3%).

5. 의의 및 주장

본 논문은 CKA 가이드 표현 정렬을 양자화된 LLM 복구를 위한 실질적인 보완책으로 자리매김합니다. 저자들은 다음과 같이 주장합니다:

기능적 정렬은 필요조건이지만 충분조건은 아니다: 출력 분포를 일치시키는 것이 심층 추론에 필요한 내부 의미론적 위상(semantic topology)의 보존을 보장하지는 않습니다.
기하 구조 보존은 저비트 일반화를 돕는다: 학생 모델이 손실 함수를 만족시키지만 내부 구조를 저해하는 "지름길" 매핑을 학습하는 것을 방지함으로써, CKA 정규화는 공격적인 양자화 과정에서도 사후 학습(특히 RL) 중에 인코딩된 미세한 의미론적 계층 구조를 유지하도록 돕습니다.
양자화 아티팩트에 대한 강건성: 포인트와이즈(pointwise) 특징 손실(예: MSE)과 달리, CKA의 스케일링 및 회전에 대한 불변성은 블록 단위 스케일링이 비정상적 분포를 유도하는 NVFP4에 매우 적합합니다.

저자들은 수학적 추론 및 코드 합성처럼 다단계 특징 전파(multi-step feature propagation)가 필요한 작업에서 내부 기하 구조를 보존하는 것이 매우 중요하며, CKA-QAD가 최소한의 엔지니어링 오버헤드로 고충실도 저비트 LLM 배포를 위한 실행 가능한 경로를 제공한다고 결론지었습니다.