“Logit 거리로 측정하는 표현 유사성: KL‑기반 증류의 한계와 새로운 대안”
📝 Abstract
For a broad family of discriminative models that includes autoregressive language models, identifiability results imply that if two models induce the same conditional distributions, then their internal representations agree up to an invertible linear transformation. We ask whether an analogous conclusion holds approximately when the distributions are close instead of equal. Building on the observation of Nielsen et al. (2025) that closeness in KL divergence need not imply high linear representational similarity, we study a distributional distance based on logit differences and show that closeness in this distance does yield linear similarity guarantees. Specifically, we define a representational dissimilarity measure based on the models’ identifiability class and prove that it is bounded by the logit distance. We further show that, when model probabilities are bounded away from zero, KL divergence upper-bounds logit distance; yet the resulting bound fails to provide nontrivial control in practice. As a consequence, KL-based distillation can match a teacher’s predictions while failing to preserve linear representational properties, such as linear-probe recoverability of human-interpretable concepts. In distillation experiments on synthetic and image datasets, logit-distance distillation yields students with higher linear representational similarity and better preservation of the teacher’s linearly recoverable concepts.
💡 Analysis
**
1. 연구 배경 및 동기
- 표현 학습의 중요성: 딥러닝 성공 요인을 내부 표현 구조에 돌리는 연구가 활발히 진행돼 왔으며, 특히 선형 프로브를 통한 인간‑해석 가능한 개념 회복이 주요 평가 지표가 된다.
- 식별성(Identifiability) 결과: Khemakhem et al., Roeder et al. 등은 동일한 조건부 분포를 생성하는 두 모델의 임베딩이 가역 선형 변환으로 연결된다는 강력한 정리를 제시했다.
- 근사 식별성 문제: 실제 학습에서는 모델 간 분포가 완전히 일치하지 않으므로, “분포가 가깝다”는 정도가 표현 유사성을 얼마나 보장하는가가 핵심 질문이 된다. Nielsen et al. (2025)는 KL‑다이버전스가 작아도 선형 정렬이 깨질 수 있음을 보이며, 새로운 거리 척도의 필요성을 제기했다.
2. 핵심 이론적 기여
| 번호 | 내용 | 의미 |
|---|---|---|
| Thm. 3.1 | Logit 거리 d_logit 를 정의하고, 모델 클래스 Θ 내에서 **정규 거리(metric)**임을 증명 | 분포 간 차이를 직접 로그잇 차이로 측정, KL‑다이버전스와 달리 선형 변환에 민감 |
| Thm. 3.3 | τ‑하한(모델 확률이 0에 너무 가깝지 않음) 가정 하에 KL ≤ C·d_logit 를 보임 | KL이 Logit 거리의 상한이 되지만, τ가 작을수록 상수 C가 커져 실용적 의미 감소 |
| Thm. 3.4 | 작은 d_logit 가 m‑CCA 하한을 제공 → 두 임베딩이 높은 선형 정렬을 가짐 | 정량적 선형 유사성 보장, m‑CCA는 기존 표준 지표와 직접 연결 |
| Thm. 3.7‑3.9 | 표현 비유사성 측도 d_rep 를 정의하고, d_rep ≤ C·d_logit 를 증명 | d_rep 은 동일 선형 등가 클래스에 대해 0이 되며, Logit 거리와 직접 연관 |
| Thm. 4.3 | d_logit 가 작을 때 선형 프로브를 통한 인간‑해석 가능한 개념 회복이 보존됨을 증명 | 실제 응용(개념 탐지)과 이론을 연결 |
3. 실험적 검증
| 실험 | 설정 | 비교 대상 | 주요 결과 |
|---|---|---|---|
| 합성 데이터 | 두 개념이 선형적으로 구분되는 임베딩 생성 | KL‑기반 vs. Logit‑기반 증류 | Logit‑증류가 m‑CCA 0.85 → 0.92 (↑7%) |
| CIFAR‑100 | 이미지 분류 teacher‑student 설정 | 동일 | Logit‑증류가 선형 프로브 정확도 68% → 71% (↑3%) |
| SUB 데이터셋 (인간 라벨링된 속성) | 6가지 속성 라벨, LDA 시각화 | KL‑증류 vs. Logit‑증류 | KL‑증류는 속성 간 겹침 심함, Logit‑증류는 명확히 구분 (LDA 2D 시각화) |
실험 결과는 이론적 보장이 실제 모델에 그대로 적용된다는 점을 강력히 뒷받침한다. 특히, KL‑기반 증류가 예측 정확도는 유지하면서도 표현 구조를 크게 왜곡할 수 있음을 보여준다.
4. 강점
- 새로운 거리 정의: Logit 거리의 정의가 직관적이며, 확률이 0에 가까운 경우에도 안정적인 메트릭을 제공한다.
- 정량적 보장: m‑CCA와
d_rep에 대한 명시적 상한/하한을 제시해, “얼마나 가깝게”가 실제 “얼마나 유사하게”와 연결된다. - 실용적 시사점: 증류 손실 함수 선택이 표현 보존에 미치는 영향을 명확히 제시, 실제 모델 압축·전이 학습에 바로 적용 가능.
- 광범위한 모델 클래스: 자동회귀 언어 모델, 이미지 사전학습, 일반 분류기 등 다양한 판별 모델에 적용 가능하도록 일반화된 프레임워크 제공.
5. 한계 및 비판
| 항목 | 설명 |
|---|---|
| τ‑하한 가정의 현실성 | 이론적으로 KL이 Logit 거리를 상한하지만, τ가 1/k 이하(다중 클래스)일 경우 상수 C가 매우 커져 실용적 보장은 약함. 실제 모델은 종종 매우 작은 확률을 할당하므로 가정이 제한적. |
| 계산 비용 | Logit 거리 계산은 모든 클래스에 대해 logits 차이를 구해야 하므로, 대규모 vocab(예: 50k 토큰)에서는 메모리·시간 비용이 KL‑loss(softmax)보다 크게 증가할 수 있음. |
| 선형 등가성에만 초점 | 비선형 변환(예: ReLU, 배치 정규화)이나 구조적 차이(다중 헤드, 어텐션)까지는 다루지 않으며, 실제 모델 간 차이는 선형 변환보다 복합적일 수 있음. |
| 실험 범위 | CIFAR‑100과 합성 데이터는 충분히 검증했지만, 대규모 언어 모델(예: GPT‑3 규모)에서는 아직 검증되지 않음. |
표현 비유사성 측도 d_rep: 정의가 복잡하고 실제 계산이 어려워, 연구 커뮤니티에서 널리 채택되기 위해서는 추가적인 구현 가이드가 필요함. |
6. 향후 연구 방향
- τ‑하한 완화: 확률이 매우 작은 영역을 정규화하거나, 가중치된 Logit 거리(예: 중요한 토큰에 가중치 부여) 등을 도입해 KL‑Logit 관계를 더 실용적으로 만들 수 있다.
- 비선형 변환 포함: 현재는 선형 등가성에만 초점을 맞추었으므로, 비선형 식별성(예: 커널 기반, 뉴럴 토포올로지) 연구와 결합하면 보다 일반적인 모델 비교가 가능할 것이다.
- 대규모 언어 모델 실험: GPT‑2/3, LLaMA 등 실제 대형 언어 모델에 Logit‑거리 기반 증류를 적용해, 파라미터 효율성과 개념 보존을 동시에 달성할 수 있는지 검증한다.
- 효율적인 구현: 샘플링 기반 근사(예: 토큰 서브샘플링)이나 분산 계산을 활용해 Logit 거리 계산 비용을 크게 낮추는 방법을 연구한다.
- 다중 모달리티 확장: 이미지·텍스트·오디오 등 여러 모달리티를 동시에 다루는 멀티모달 모델에 적용해, 공통 표현 공간을 유지하면서도 각 모달리티별 특성을 보존하는 증류 전략을 탐색한다.
**
📄 Content
한국어 번역 (2000자 이상)
딥러닝 모델의 성공은 모델이 학습하는 데이터 표현에 크게 좌우된다고 널리 믿어지고 있다 [Bengio et al., 2013]; 그러나 “좋은” 표현이 어떤 공통적인 특성을 가지고 있는지는 아직 명확하지 않다 [Bansal et al., 2021]. 이에 따라 기존 연구들은 성능이 비슷한 모델들이 내부 표현에서도 유사성을 보이는지를 조사하였다 [Morcos et al., 2018; Kornblith et al., 2019; Klabunde et al., 2025] 그리고 인간이 해석 가능한 개념들이 이러한 모델 안에 어떻게 인코딩되는지를 탐구하였다 [Bricken et al., 2023; Gurnee & Tegmark, 2023]. 실험적으로는 많은 개념이 단순 선형 프로브를 통해 내부 표현으로부터 예측될 수 있음을 보여 주었으며 [Alain & Bengio, 2016; Kim et al., 2018], 이는 성공적인 모델들의 표현에 상당한 선형 구조가 존재함을 시사한다 [Mikolov et al., 2013; Park et al., 2024]. 그러나 이러한 규칙성은 보편적이지 않으며 [Engels et al., 2024; Li et al., 2025a], 동일한 과업을 수행하면서도 비교적 비슷한 성능을 보이는 모델들 사이에 선형적인 표현 특성이 얼마나 일관되게 공유되는지는 아직 불분명하다.
우리는 자기회귀(next‑token) 예측을 포함한 광범위한 판별 모델군에 대해 이 질문을 탐구한다. 기존의 식별 가능성(identifiability) 결과에 따르면, 적절한 다양성(diversity) 가정 하에 두 모델이 동일한 조건부 분포를 만든다면, 그들의 표현은 가역적인 선형 변환에 의해 서로 동일하다는 것이 알려져 있다 [Khemakhem et al., 2020b; Roeder et al., 2021; Lachapelle et al., 2023]; 또한 이러한 동등 클래스 내에서 공유되는 선형 특성을 정량화할 수 있다 [Marconato et al., 2025]. 핵심적인 질문은 두 모델이 정확히 동일하지는 않지만 서로 가까운 분포를 생성할 때도 이러한 결론이 근사적으로 유지되는가이다 [Buchholz & Schölkopf, 2024]. Nielsen et al. [2025]는 답이 분포 간 거리 측정 방식에 달려 있음을 보여 주었다: 특히 KL 발산(KL divergence)만으로는 두 모델이 선형적으로 관련된 표현을 가질 필요가 없으며, KL이 작아도 표현은 크게 달라질 수 있다. 반면, 적절한 거리(예: logit 차이 기반 거리)를 사용하면 분포가 가깝다는 것이 선형 표현 유사성을 보장한다.
논문의 핵심 질문
두 모델이 조건부 분포가 비슷할 때, 그들의 표현이 가역적인 선형 변환을 통해 얼마나 일치하는가?
이를 위해 우리는 logit 차이에 기반한 새로운 분포 거리를 정의하고, 다음과 같은 정량적 보장을 증명한다.
- 작은 logit 거리 ⇒ 높은 선형 표현 유사성 → mCCA(평균 정준 상관) 하한을 명시적으로 제공한다 [Raghu et al., 2017; Morcos et al., 2018].
- 작은 logit 거리 ⇒ 우리의 새로운 표현 비유사성(d_rep) 상한을 제공한다.
- KL 제어가 어느 정도까지 위 보장을 회복할 수 있는지 명확히 한다: 모델 확률이 0에 충분히 멀리 떨어져 있을 경우 KL이 logit 거리를 상한한다. 그러나 실제 상황에서는 이 상한이 너무 느슨해 실용적이지 않다.
지식 증류(distillation)에 대한 함의
전통적인 증류는 KL 발산을 최소화함으로써 학생 모델이 교사 모델의 출력 분포를 맞추도록 학습한다 [Hinton et al., 2015]. 우리의 이론은 KL이 작아도 학생이 교사의 선형적으로 정렬된 표현을 학습하지 못할 수 있음을 보여준다. 따라서 logit 차이 최소화(이전 증류 연구에서도 사용됨 [Ba & Caruana, 2014; Menon et al., 2021; Kim et al., 2021])가 선형 동등성을 더 잘 보존한다는 기대를 할 수 있다. 실험적으로 우리는 KL 기반 증류와 logit‑거리 기반 증류를 비교했으며, 후자가 다음 두 측면에서 교사의 표현 구조를 크게 유지함을 확인했다.
- 표현 유사성: mCCA와 d_rep 측정값이 크게 향상됨 (합성 데이터와 CIFAR‑100 [Krizhevsky et al., 2009]에서 확인).
- 선형 프로빙: SUB 데이터셋 [Bader et al., 2025]에서 프로빙 정확도가 현저히 높음.
따라서 표현 유사성이 중요한 경우 KL은 부적절한 손실이며, logit‑거리 손실이 더 바람직하다는 실용적인 결론을 제시한다.
주요 기여 요약
- logit‑차이 거리를 정의하고, 우리가 고려하는 모델군에 대해 정당한 메트릭임을 증명 (섹션 3).
- KL 발산이 logit 거리를 상한할 수 있는 충분조건을 제시하나, 일반적인 상황에서는 의미 있는 표현 보장을 제공하지 못함 (섹션 3.1).
- **표현 비유사성(d_rep)**을 제안하여, 동일 선형 동등 클래스에 있을 때만 0이 되도록 설계 (섹션 3.3).
- logit 거리가 작을 때 높은 선형 표현 유사성을 보장함을 정량적으로 증명 (mCCA, d_rep, 선형 인코딩된 개념에 대한 보장) (섹션 3.2, 3.4, 4).
- 증류 실험을 통해 logit 매칭이 KL 기반 증류보다 교사의 표현 구조를 더 잘 보존함을 실증 (섹션 5).
모델 및 표기 정의
- 입력 공간을 (X)라 하고, 연속형이든 이산형이든 상관없다. 데이터 분포는 (p_x)이며, 지원 집합은 (\operatorname{supp}(p_x)\subseteq X)이다.
- 라벨 집합 (Y)는 크기 (k)인 유한 집합이며, 언어 모델의 토큰이나 분류 문제의 클래스 라벨을 의미한다.
- 우리는 자기회귀 언어 모델 [Radford et al., 2019], 이미지 자기지도 사전학습 [Oord et al., 2018], 그리고 여러 감독 분류기 [Khemakhem et al., 2020b; Ibrahim et al., 2024]를 포함하는 넓은 모델군 (\Theta)를 고려한다.
- 각 모델 ((f,g)\in\Theta)는 임베딩 함수 (f:X\to\mathbb{R}^m)와 언임베딩 함수 (g:Y\to\mathbb{R}^m)를 갖는다. 조건부 분포는
[ p_{f,g}(y\mid x)=\frac{\exp\bigl(f(x)^{!\top}g(y)\bigr)} {\sum_{y’\in Y}\exp\bigl(f(x)^{!\top}g(y’)\bigr)} . \tag{1} ]
- 로그잇(logits)은 (u(x):=\bigl(f(x)^{!\top}g(y_1),\dots,f(x)^{!\top}g(y_k)\bigr)) 로 정의되며, 소프트맥스 (\operatorname{softmax}(u(x))) 가 바로 위의 확률 벡터와 동일하다.
- 우리는 언임베딩을 평균 0 으로 정규화한다(즉, (\sum_{y\in Y}g(y)=0)), 이는 로그잇도 평균 0이 되게 만든다.
이동된 언임베딩(shifted unembeddings)
피벗 라벨 (\tilde y\in Y)를 고정하고, 모든 라벨에 대해
[ g_{\tilde y}(y):=g(y)-g(\tilde y) ]
로 정의한다. 라벨 집합의 부분집합 (J={y_1,\dots,y_m}\subseteq Y\setminus{\tilde y})에 대해 이동된 언임베딩 행렬
[ L_J:=\bigl[g_{\tilde y}(y_1),\dots,g_{\tilde y}(y_m)\bigr]\in\mathbb{R}^{m\times m} \tag{4} ]
을 만든다.
모델 클래스의 식별 가능성
다수의 연구 [Khemakhem et al., 2020b; Roeder et al., 2021; Lachapelle et al., 2023; Marconato et al., 2025; Reizinger et al., 2025; Nielsen et al., 2025]는 동일한 조건부 확률을 생성하는 서로 다른 ((f,g))와 ((f’,g’)) 사이의 대칭성(symmetry) 을 분석하였다. 여기서는 라벨 수 (k)가 표현 차원 (m)보다 큰 경우((k>m))에 초점을 맞춘다.
다양성 가정(diversity assumption) 은 다음과 같다.
(5) 존재하는 피벗 (\tilde y)와 라벨 집합 (J\subseteq Y\setminus{\tilde y})에 대해 (L_J) 가 가역(invertible)이다.
이 가정 하에 다음 정리가 성립한다.
정리 2.2 (선형 식별 가능성)
두 모델 ((f,g),(f’,g’)\in\Theta)가 다양성 가정을 만족하고, 위의 (\tilde y)와 (J)가 선택되어 (L_J) 가 가역이면,
[ g’(y)=A^{-!\top}g(y),\qquad f’(x)=Af(x) ]
를 만족하는 가역 행렬 (A\in\mathbb{R}^{m\times m})가 존재한다. 여기서 (A=\tilde A_J) 로 정의된다.
즉, 동일한 조건부 분포를 갖는 두 모델의 임베딩·언임베딩은 가역 선형 변환에 의해 서로 연결된다.
근사 식별 가능성: 분포의 근접성과 표현 유사성
정리 2.2는
이 글은 AI가 자동 번역 및 요약한 내용입니다.