논문 리뷰
논문 리뷰: EMO - Embedding Model Distillation via Intra-Model Relation and Optimal Transport Alignments
EMO: Embedding Model Distillation via Intra-Model Relation and Optimal Transport Alignments 리뷰
서로 다른 토크나이저를 쓰는 대형 임베딩 모델에서 소형 모델로, 어텐션 구조와 표현을 동시에 증류하는 프레임워크
논문 정보
| 항목 | 내용 |
|---|---|
| 제목 | EMO: Embedding Model Distillation via Intra-Model Relation and Optimal Transport Alignments |
| 저자 | Minh-Phuc Truong, Hai An Vu, Tu Vu, Diep Thi-Ngoc Nguyen, Linh Ngo Van, Thien Huu Nguyen, Trung Le (HUST, ByteDance, VNU, U of Oregon, Monash) |
| 학회/저널 | EMNLP 2025 (Main Conference) |
| 논문 링크 | ACL Anthology |
1. 들어가며
LLM 기반 텍스트 임베딩 모델이 MTEB 리더보드를 휩쓸고 있지만, 실제 서비스에 7B 파라미터짜리 모델을 띄우는 건 만만치 않은 일이다. RAG 파이프라인에 임베딩 모델을 붙이려 해도, 추론 비용과 레이턴시가 발목을 잡는다. 그래서 지식 증류(Knowledge Distillation, KD)가 자연스러운 해법으로 떠오르는데, 여기에 한 가지 까다로운 제약이 있다.
DistilBERT, TinyBERT, MiniLM 같은 기존 증류 방법은 teacher와 student가 같은 토크나이저를 쓴다고 가정한다. 토큰 위치가 1:1로 대응되니까 logit이든 hidden state든 attention이든 직접 비교할 수 있는 거다. 그런데 실제로는 teacher가 Mistral 7B(SentencePiece)이고 student가 BERT(WordPiece)인 식으로, 토크나이저 자체가 다른 경우가 대부분이다. "watching"이라는 단어를 teacher는 ["watch", "ing"]으로, student는 ["watching"]으로 토큰화하면, 시퀀스 길이도 다르고 토큰 경계도 안 맞는다.
Cross-Tokenizer KD(CTKD) 연구들이 이 문제를 다루고 있지만, ULD, MinED, DSKD, MultilevelOT 등 기존 방법은 대부분 출력 logit 정렬에 집중한다. 모델 내부의 어텐션 패턴이나 hidden state에 담긴 구조적 지식은 제대로 전달하지 못하는 셈이다. 같은 토크나이저를 쓰는 KD에서는 이미 TinyBERT, MiniLM이 어텐션 행렬 증류의 효과를 입증했건만, 크로스 토크나이저 상황에서는 이를 적용하기가 어려웠다.
이 논문은 바로 이 빈틈을 파고든다. **EMO(Embedding Model Distillation via Intra-Model Relation and Optimal Transport Alignments)**라는 프레임워크를 제안하여, 서로 다른 토크나이저 사이에서도 어텐션 구조와 토큰 표현을 동시에 증류할 수 있는 방법을 보여준다.
2. 기존 연구의 한계
2.1 같은 토크나이저 기반 증류
KD의 발전 경로를 보면, 초기 logit 매칭(Hinton et al., 2015)에서 시작해 hidden state 증류(PKD, Sun et al., 2019), 어텐션 행렬 증류(TinyBERT, Jiao et al., 2019; MiniLM, Wang et al., 2020)로 진화해왔다. 이 과정에서 "모델 내부의 관계적 지식(relational knowledge)"이 증류에 매우 중요하다는 것이 반복적으로 확인됐다. 어텐션 행렬은 토큰 간 의존 관계와 맥락 구조를 담고 있어서, 이를 전달하면 student 모델의 품질이 크게 올라간다.
문제는 이 방법들이 모두 teacher-student 간 토큰 위치가 정확히 일치한다는 전제 위에 서 있다는 것이다.
2.2 크로스 토크나이저 증류(CTKD)
토크나이저가 다른 상황을 다루는 연구도 여러 갈래로 진행됐다:
- ULD (Boizard et al., 2025): Optimal Transport로 출력 logit 분포를 정렬
- MinED (Wan et al., 2024): 최소 편집 거리(Minimum Edit Distance) 기반 동적 프로그래밍으로 토큰 시퀀스를 정렬
- DSKD (Zhang et al., 2024): 프로젝션과 크로스 모델 어텐션으로 출력 공간을 통합
- MultilevelOT (Cui et al., 2024): 다층 Optimal Transport 전략
이 방법들의 공통적 한계는 출력 레벨 정렬에 머문다는 점이다. 모델 내부의 어텐션 패턴이나 중간 레이어의 표현 구조 — 같은 토크나이저 KD에서 이미 효과가 입증된 지식 원천 — 을 활용하지 못한다. 특히 텍스트 임베딩 모델처럼 맥락 이해와 구조적 파악이 중요한 도메인에서, 이 정보 손실은 성능 병목이 된다.
또 하나 간과하기 쉬운 문제가 있다. STS(Semantic Textual Similarity)처럼 연속 스칼라를 예측하는 회귀 태스크에서는 logit 벡터 자체가 존재하지 않으므로, logit 기반 CTKD 방법(ULD, MinED, MultilevelOT)을 적용할 수가 없다. 이건 실용적으로 꽤 큰 제약이다.
3. EMO의 접근: 두 가지 상보적 증류
EMO의 아이디어는 명쾌하다. 크로스 토크나이저 환경에서도 모델 내부 지식을 전달하되, **관계적 지식(어텐션 구조)**과 **표현적 지식(hidden state)**을 각각 다른 도구로 정렬하자는 것이다.
첫째, IRA(Intra-Model Relational Alignment) 는 MinED로 매핑된 토큰 쌍 중 가장 중요한 top-개에 대해 어텐션 행렬을 CKA(Centered Kernel Alignment)로 정렬한다. 이건 teacher 내부에서 "어떤 토큰이 어떤 토큰에 주목하는가"라는 관계 구조를 student에 심어주는 역할이다.
둘째, OTIS(Optimal Transport with Importance-Scored Mass Assignment) 는 마지막 hidden state를 Optimal Transport로 직접 정렬하되, 어텐션 기반 중요도 점수로 각 토큰의 mass를 차등 배정한다. 중요한 토큰의 표현은 더 무겁게, 덜 중요한 토큰은 가볍게 취급하는 것이다.
IRA가 "내부 구조 보존"을, OTIS가 "모델 간 표현 정렬"을 담당하면서 상호 보완적으로 작동한다. 그리고 둘 다 logit에 의존하지 않기 때문에 분류든 회귀든 태스크를 가리지 않는다.
4. 제안 방법 (Method)
4.1 전체 구조
EMO 프레임워크의 전체 워크플로우는 다음과 같다:
- 입력 텍스트를 teacher와 student 토크나이저로 각각 토큰화
- MinED로 1:1 매핑 가능한 토큰 쌍 개를 식별
- Teacher의 마지막 레이어 어텐션에서 중요도 점수를 계산하여 top-개 토큰 선정
- IRA: 선정된 토큰의 어텐션 패턴을 CKA로 정렬 (마지막 개 레이어)
- OTIS: 마지막 hidden state를 중요도 기반 mass를 가진 Optimal Transport로 정렬
- 태스크 손실과 증류 손실을 결합하여 student 학습

Figure 1: EMO 프레임워크의 전체 워크플로우. MinED로 토큰을 매핑한 뒤, 왼쪽의 IRA는 어텐션 행렬 간 CKA 정렬을, 오른쪽의 OTIS는 중요도 기반 mass를 활용한 Optimal Transport로 hidden state 정렬을 수행한다. (원논문)
Figure 1을 보면 "I'm watching funny videos"라는 예시로 전체 흐름이 잘 드러난다. Teacher는 이를 ["I'm", "watch", "ing", "fun", "ny", "videos"]로, student는 ["I'm", "watching", "funny", "video", "s"]로 토큰화한다. MinED가 대응 관계를 찾고, 어텐션 행렬에서 중요도 점수를 뽑아 top- 토큰을 선별한다. 왼쪽 경로(IRA)에서는 이 토큰들의 sub-attention 행렬에 CKA를 적용하고, 오른쪽 경로(OTIS)에서는 중요도 점수를 mass로 변환하여 Sinkhorn 알고리즘 기반 OT를 수행한다. 두 손실이 합산되어 student를 학습시키는 구조다.
4.2 MinED 기반 토큰 매핑
서로 다른 토크나이저의 토큰 시퀀스를 정렬하는 첫 단계로, Minimum Edit Distance(MinED)를 사용하여 개의 1:1 매핑 토큰 쌍을 찾는다. Student 시퀀스 (길이 )와 teacher 시퀀스 (길이 )에 대해 인 매핑 쌍 를 구한다.
실제로 Bert-base-uncased(student)와 LLM2Vec-Mistral 7B(teacher) 사이에서 MinED 매핑이 얼마나 잘 작동하는지는 뒤의 Appendix A에서 다루겠지만, 미리 말하면 평균 약 75%의 토큰이 1:1 매핑된다. 상당히 높은 비율이다.
4.3 중요 토큰 식별
매핑된 개 토큰 중에서도 모든 토큰이 동등하게 중요하지는 않다. Li et al. (2023)의 관찰에서 영감을 받아, teacher 모델의 마지막 레이어 어텐션에서 각 토큰의 중요도를 계산한다. 마지막 레이어는 토큰 중요도에 대해 더 전역적(global) 관점을 제공하기 때문이다.
Teacher의 마지막 레이어에서 모든 헤드를 평균한 어텐션 행렬 를 구하고, 번째 토큰의 중요도 점수를 다음과 같이 계산한다:
각 토큰이 다른 모든 토큰으로부터 받는 어텐션의 총합이 곧 중요도가 되는 것이다. 직관적으로, 많은 토큰이 주목하는 토큰일수록 문장의 의미 파악에 중요하다고 볼 수 있다.
매핑된 토큰 인덱스 중에서 이 점수가 높은 상위 개를 선별하여 으로 정의하고, 대응하는 student 인덱스를 으로 둔다. 의 선택은 Section 5에서 분석하는데, 이 가장 좋은 성능을 보인다. 모든 매핑 토큰이 아닌 상위 일부만 쓰는 게 노이즈 감소 효과를 가져오는 셈이다.
4.4 IRA: 어텐션 행렬의 CKA 정렬
IRA의 목표는 teacher와 student의 어텐션 패턴 구조를 맞추는 것이다. 여기서 CKA(Centered Kernel Alignment)가 등장한다.
CKA 배경
CKA는 Kornblith et al. (2019)이 제안한 표현 유사도 측정 도구로, CCA(Canonical Correlation Analysis)의 한계를 극복한다. CCA가 등방성 스케일링이나 회전에 민감한 반면, CKA는 이에 불변(invariant)하다. CKA의 핵심은 개별 특징을 직접 비교하는 대신, 각 표현 공간 내에서 학습된 유사도 구조를 비교한다는 점이다.
CKA는 HSIC(Hilbert-Schmidt Independence Criterion)를 기반으로 한다. 두 표현 집합 , 에 대해 커널 행렬 , 을 구성하고, 센터링 행렬 로 평균을 제거한 뒤 의존성을 측정한다:
CKA는 HSIC를 정규화하여 0과 1 사이 값을 갖는다:
어텐션 행렬에 CKA 적용
Student의 번째 레이어는 teacher의 번째 레이어에 매핑된다. 각 모델에서 top- 토큰의 어텐션 행을 추출하여 sub-attention 행렬을 구성한다:
- Student: — 에 해당하는 행
- Teacher: — 에 해당하는 행
선형 커널을 사용하여 커널 행렬을 정의한다:
센터링된 행렬은:
HSIC는 다음과 같이 계산된다:
이를 정규화하면 선형 CKA가 된다:
여기서 , 이다. CKA 값은 [0, 1] 범위에 있으며, 1에 가까울수록 두 어텐션 구조가 유사하다.
최종 IRA 손실은 student의 마지막 개 레이어에 걸쳐 적용된다:
를 쓰는 이유는, CKA가 1일 때(완벽한 정렬) 손실이 0이 되도록 하기 위함이다. 이전 연구(Dasgupta and Cohn, 2025)에서는 CKA를 hidden state 비교에 사용했지만, EMO는 이를 어텐션 행렬에 직접 적용한다는 점이 차별화된다. 는 top- 중요 토큰이 자신의 전체 시퀀스에 대해 보이는 어텐션 패턴의 구조적 유사성을 측정하는 것이다.
중요한 포인트는, sub-attention 행렬의 차원이 teacher()와 student()에서 열 수가 다르다는 것이다. 그런데 CKA는 행(토큰) 수만 같으면 열 수가 달라도 비교가 가능하다. 커널 행렬이 으로 동일하기 때문이다. 이 성질 덕분에 시퀀스 길이가 달라도 어텐션 구조를 비교할 수 있다.
4.5 OTIS: 중요도 기반 Optimal Transport 정렬
IRA가 어텐션의 관계적 구조를 전달한다면, OTIS는 마지막 hidden state의 토큰 표현 자체를 직접 정렬한다. Optimal Transport를 사용하되, 모든 토큰을 균등하게 취급하지 않고 어텐션 기반 중요도로 차등을 둔다.
중요도 기반 Mass 배정
Teacher 토큰의 mass는 Eq.6의 중요도 점수를 정규화하여 할당한다:
Teacher의 마지막 hidden state 에 대한 경험적 분포는:
Student에는 매핑된 토큰 쌍 을 통해 teacher의 중요도를 전달한다. 매핑된 student 토큰 에는 대응하는 teacher 토큰의 mass를 할당하고(), 매핑되지 않은 토큰에는 teacher 토큰 중 최소 mass를 배정한다(). 이후 전체 합이 1이 되도록 정규화한다:
이 설계의 의도는 분명하다. 의미적으로 중요한 토큰의 표현 정렬에 더 많은 가중치를 두고, 매핑되지 않았거나 덜 중요한 토큰은 가볍게 처리하겠다는 것이다.
비용 행렬 계산
Teacher의 hidden state를 학습 가능한 매핑 행렬 로 student 공간에 프로젝션한 뒤, 유사도 행렬을 계산한다:
로 나누는 건 수치 안정성을 위한 스케일링이다. 유사도 행렬을 행 단위 softmax로 정규화한 뒤, 비용 행렬을 유도한다:
유사도가 높으면 비용이 낮고, 유사도가 낮으면 비용이 높아지는 직관적인 구조다.
OT 정렬 손실
엔트로피 정규화된 OT 문제를 Sinkhorn 알고리즘으로 풀어 최적 transport plan 를 구한다:
여기서 는 엔트로피 정규화 항이고, 은 정규화 강도를 제어한다. 최종 OT 손실은:
OTIS는 마지막 레이어에만 적용된다:
OT의 계산 비용이 상당하기 때문에 마지막 레이어로 제한한 것인데, Appendix C의 분석에서 이 선택이 효율과 효과 사이의 합리적 균형점임을 확인할 수 있다.
4.6 전체 증류 손실
최종 학습 목적 함수는 태스크 손실과 증류 손실의 가중 합이다:
이 표준 cross-entropy 손실 와 증류 손실() 사이의 비중을 조절한다. 대부분의 데이터셋에서 가 최적이었다(Appendix B, Table 8 참조).
5. 실험 결과
5.1 실험 설정
모델 구성:
- Teacher: LLM2Vec Mistral 7B SFT — MTEB 리더보드의 SOTA 임베딩 모델
- Student: Bert-base-uncased (110M 파라미터)
- 파라미터 규모 차이가 약 63배로, 상당히 공격적인 압축 시나리오다
태스크: 세 종류의 임베딩 기반 태스크를 평가한다:
- 텍스트 분류: Patent, Imdb, Banking77
- 문장 쌍 분류: SciTail, ConTRoL-NLI, Anli_r2
- 의미 유사도(STS): STSB, STS12, SICK
평가 지표: 분류 태스크는 Accuracy/Precision/Recall, STS는 Spearman 상관계수()
학습 설정: Teacher는 LoRA(rank=256, alpha=32)로 fine-tune하고, student는 전체 파라미터를 fine-tune한다. 에폭 5, 학습률 , 배치 크기 4, cosine 스케줄러를 사용한다.
| Dataset | Train | Validation | Test |
|---|---|---|---|
| Patent | 25,000 | 5,000 | 5,000 |
| Imdb | 25,000 | - | 25,000 |
| Banking77 | 10,000 | - | 3,080 |
| SciTail | 23,100 | 1,300 | 2,130 |
| ConTRoL-nli | 6,720 | 799 | 805 |
| Anli_r2 | 45,500 | 1,000 | 1,000 |
| STSB | 5,750 | 1,500 | 1,380 |
| STS12 | 2,230 | - | 3,110 |
| SICK | 4,500 | 500 | 4,823 |
Table 9 (원논문): 데이터셋 통계. 다양한 규모와 도메인의 9개 데이터셋을 사용한다.
데이터셋 규모가 수천에서 수만 샘플로 다양하고, Patent이나 Anli_r2처럼 비교적 큰 데이터셋과 ConTRoL-nli처럼 작은 데이터셋이 섞여 있어 다양한 조건에서의 성능을 확인할 수 있다.
5.2 주요 결과: 분류 및 문장 쌍 분류
| Method | Patent | SciTail | ||||
|---|---|---|---|---|---|---|
| Acc | Prec | Recall | Acc | Prec | Recall | |
| LLM2Vec Mistral 7B SFT (Teacher) | 70.0 | 67.7 | 66.1 | 96.1 | 96.0 | 95.8 |
| Bert SFT (Student) | 63.1 | 58.7 | 54.4 | 88.1 | 87.7 | 88.8 |
| ULD | 64.8 | 61.4 | 60.9 | 87.0 | 86.4 | 87.8 |
| DSKD | 64.0 | 60.0 | 58.8 | 88.0 | 87.3 | 88.8 |
| MinED | 65.0 | 61.6 | 60.8 | 86.9 | 86.1 | 87.5 |
| MultilevelOT | 64.6 | 60.4 | 59.0 | 88.2 | 88.0 | 89.1 |
| EMO | 66.5 | 63.3 | 62.4 | 90.9 | 90.1 | 91.2 |
| Method | Imdb | ConTRoL-nli | ||||
|---|---|---|---|---|---|---|
| Acc | Prec | Recall | Acc | Prec | Recall | |
| LLM2Vec Mistral 7B SFT (Teacher) | 96.6 | 96.6 | 96.6 | 63.6 | 62.7 | 62.6 |
| Bert SFT (Student) | 91.3 | 91.4 | 91.3 | 42.1 | 38.6 | 37.5 |
| ULD | 92.5 | 92.6 | 92.5 | 45.4 | 45.3 | 45.3 |
| DSKD | 93.4 | 93.5 | 93.4 | 42.2 | 41.2 | 39.7 |
| MinED | 92.5 | 92.5 | 92.5 | 47.1 | 47.0 | 47.2 |
| MultilevelOT | 93.3 | 93.4 | 93.3 | 42.5 | 41.4 | 40.1 |
| EMO | 94.2 | 94.3 | 94.2 | 48.6 | 48.2 | 48.1 |
| Method | Banking77 | Anli_r2 | ||||
|---|---|---|---|---|---|---|
| Acc | Prec | Recall | Acc | Prec | Recall | |
| LLM2Vec Mistral 7B SFT (Teacher) | 93.3 | 93.5 | 93.3 | 67.1 | 67.8 | 67.0 |
| Bert SFT (Student) | 85.7 | 86.4 | 85.7 | 42.7 | 42.6 | 42.6 |
| ULD | 91.4 | 91.9 | 91.4 | 44.8 | 44.7 | 44.7 |
| DSKD | 91.2 | 91.7 | 91.2 | 43.1 | 43.4 | 43.0 |
| MinED | 90.0 | 91.2 | 90.0 | 46.4 | 46.6 | 46.4 |
| MultilevelOT | 89.4 | 90.4 | 89.4 | 44.1 | 44.1 | 43.9 |
| EMO | 92.3 | 92.7 | 92.3 | 47.6 | 47.8 | 47.5 |
Table 1 (원논문): 분류 및 문장 쌍 분류 태스크 결과. EMO가 모든 데이터셋에서 CTKD 베이스라인을 상회한다.
EMO가 6개 데이터셋 전체에서 모든 CTKD 베이스라인을 능가한다. 특히 눈에 띄는 건 SciTail에서의 성능이다. Student baseline이 88.1%인데, 기존 CTKD 방법들은 86.9~88.2%로 오히려 성능이 떨어지거나 미미하게 오르는 데 그치는 반면, EMO는 90.9%로 2.8%p 향상시킨다. teacher(96.1%)와의 간극을 상당 부분 좁힌 것이다.
ConTRoL-nli도 흥미로운데, 이 데이터셋은 teacher 자체가 63.6%밖에 안 되는 어려운 태스크다. Student baseline은 42.1%로 처참하고, 대부분의 CTKD 방법도 42~47%대에 머무는데, EMO는 48.6%로 가장 높다. MinED(47.1%)와 비교하면 의미 있는 차이이고, DSKD(42.2%)나 MultilevelOT(42.5%)와 비교하면 격차가 크다.
Anli_r2에서도 비슷한 양상이다. 다만 여기서는 MinED가 46.4%로 의외로 선전하는데, 이는 MinED의 토큰 매핑이 NLI 태스크에서 효과적으로 작동함을 시사한다. EMO도 MinED의 토큰 매핑을 기반으로 하므로, 이 장점을 물려받으면서 추가적인 구조 정렬로 성능을 더 끌어올린 것으로 보인다.
5.3 주요 결과: 의미 유사도(STS)
| Dataset | Method | Spearman Corr () |
|---|---|---|
| STS-B | LLM2Vec Mistral 7B SFT (Teacher) | 90.8 |
| Bert SFT (Student) | 75.1 | |
| DSKD | 78.3 | |
| EMO | 81.3 | |
| STS12 | LLM2Vec Mistral 7B SFT (Teacher) | 80.42 |
| Bert SFT (Student) | 49.7 | |
| DSKD | 65.3 | |
| EMO | 75.3 | |
| SICK | LLM2Vec Mistral 7B SFT (Teacher) | 88.9 |
| Bert SFT (Student) | 61.1 | |
| DSKD | 78.7 | |
| EMO | 80.1 |
Table 2 (원논문): STS 태스크 결과 (Spearman ). STS는 회귀 태스크이므로 logit 기반 방법(ULD, MinED, MultilevelOT)은 적용 불가하여 DSKD만 비교 대상이다.
STS 결과에서 EMO의 강점이 더 극적으로 드러난다. STS12에서 student baseline이 49.7에 불과한데, DSKD가 65.3으로 올리고 EMO는 75.3까지 끌어올린다. DSKD 대비 10.0%p, baseline 대비 25.6%p 향상이다. Teacher(80.42)와의 차이가 5.12%p밖에 안 되는데, 파라미터 수 차이가 63배라는 걸 감안하면 인상적인 수치다.
STS-B에서는 EMO(81.3) vs DSKD(78.3)로 3.0%p 차이, SICK에서는 80.1 vs 78.7로 1.4%p 차이다. 격차가 STS12보다 작은 이유는, STS-B와 SICK이 상대적으로 "쉬운" 데이터셋이라 baseline부터 높기 때문으로 해석된다.
또한 ULD, MinED, MultilevelOT가 STS에 적용조차 안 된다는 점에서, EMO의 태스크 범용성이 돋보인다. Logit에 의존하지 않는 설계 덕분에 분류와 회귀를 모두 커버하는 것이다.
5.4 Ablation Study
| Dataset | Method | Accuracy | Precision | Recall | Spearman |
|---|---|---|---|---|---|
| Patent | Bert SFT | 63.1 | 58.7 | 54.4 | - |
| EMO w/o OTIS | 65.3 | 61.8 | 60.9 | - | |
| EMO w/o IRA | 64.7 | 59.8 | 59.2 | - | |
| EMO | 66.5 | 63.3 | 62.4 | - | |
| SciTail | Bert SFT | 88.1 | 87.7 | 88.8 | - |
| EMO w/o OTIS | 88.8 | 88.1 | 89.5 | - | |
| EMO w/o IRA | 87.2 | 86.9 | 88.3 | - | |
| EMO | 90.9 | 90.1 | 91.2 | - | |
| STSB | Bert SFT | - | - | - | 75.1 |
| EMO w/o OTIS | - | - | - | 80.9 | |
| EMO w/o IRA | - | - | - | 78.5 | |
| EMO | - | - | - | 81.3 |
Table 3 (원논문): IRA와 OTIS 각 구성요소의 기여도를 보여주는 ablation 결과.
IRA만 사용(EMO w/o OTIS)해도, OTIS만 사용(EMO w/o IRA)해도, 모두 student baseline을 넘는다. 각각이 독립적으로 효과가 있다는 뜻이다. 그런데 IRA의 기여가 상대적으로 더 크다는 게 흥미롭다. Patent에서 IRA만(65.3) vs OTIS만(64.7), STSB에서 IRA만(80.9) vs OTIS만(78.5). 어텐션 구조 정렬이 hidden state 정렬보다 조금 더 효과적인 셈인데, 이건 어텐션 패턴이 담고 있는 관계적 정보가 임베딩 품질에 더 직접적으로 기여하기 때문일 수 있다.
SciTail에서는 약간 다른 양상이 나타난다. EMO w/o IRA(87.2)가 오히려 baseline(88.1)보다 낮다. OTIS만으로는 이 데이터셋에서 오히려 방해가 될 수 있다는 뜻인데, IRA와 결합하면 90.9%로 크게 뛰어오른다. 두 구성요소의 상호 보완적 관계를 잘 보여주는 결과다.
5.5 Top- 토큰 선택의 영향
| Dataset | DSKD (Baseline) | |||
|---|---|---|---|---|
| STSB | 80.2 | 80.6 | 81.3 | 78.3 |
| SICK | 76.9 | 79.4 | 80.1 | 78.7 |
Table 4 (원논문): CKA 어텐션 증류에 사용할 top- 토큰 수의 영향 (STS 태스크, Spearman ). 은 전체 1:1 매핑 토큰 수.
(매핑된 모든 토큰 사용)일 때보다 (상위 1/3만 사용)일 때 성능이 일관되게 더 좋다. STSB에서 80.2 → 81.3, SICK에서 76.9 → 80.1. 특히 SICK에서 과 의 차이가 3.2%p로 꽤 크다.
이 결과의 함의는 분명하다. 매핑된 토큰 중에도 어텐션 구조 전달에 유용한 토큰과 그렇지 않은 토큰이 있고, 덜 중요한 토큰을 포함하면 CKA 계산에 노이즈가 유입된다는 것이다. "not all tokens are what you need"라는 Rho-1(Lin et al., 2025)의 관찰과도 맥이 닿는 결과다. 어떤 값을 써도 DSKD baseline(78.3, 78.7)은 넘기 때문에, top- 선택이 필수는 아니지만 성능 최적화에 도움이 된다.
5.6 Teacher 모델 변경에 대한 견고성
| Method | Accuracy | Precision | Recall |
|---|---|---|---|
| BGE SFT (Teacher) | 94.3 | 94.1 | 93.9 |
| Bert SFT (Student) | 88.1 | 87.7 | 88.8 |
| ULD | 91.2 | 90.6 | 91.9 |
| DSKD | 91.5 | 91.2 | 91.0 |
| MinED | 89.5 | 88.9 | 90.4 |
| MultilevelOT | 91.4 | 91.2 | 90.9 |
| EMO | 92.7 | 91.8 | 92.3 |
Table 5 (원논문): SciTail에서 BGE-M3를 teacher로 사용한 결과. Teacher를 바꿔도 EMO가 모든 CTKD 방법을 상회한다.
LLM2Vec Mistral 7B 대신 BGE-M3(Chen et al., 2024)를 teacher로 사용해도 EMO가 92.7%로 가장 높다. ULD(91.2%), DSKD(91.5%), MultilevelOT(91.4%) 모두 91%대에 그치는 반면, EMO만 92%를 넘는다. 이건 EMO가 특정 teacher 모델에 과적합된 것이 아니라, 다양한 teacher에 범용적으로 적용 가능하다는 것을 보여준다.
한 가지 주목할 점은, BGE teacher(94.3%)가 LLM2Vec teacher(96.1%)보다 SciTail 성능이 낮은데, EMO의 student 성능도 이에 비례하여 낮아진다는 것이다(LLM2Vec 기반 90.9% vs BGE 기반 92.7%). 다만 BGE 기반에서 teacher-student gap(94.3 - 92.7 = 1.6%p)이 LLM2Vec 기반(96.1 - 90.9 = 5.2%p)보다 훨씬 작다. teacher 성능이 낮을수록 student가 더 가까이 따라갈 수 있다는, 직관적으로 납득이 가는 패턴이다.
6. 추가 실험 및 부록
6.1 Appendix A: 토큰 매핑 비율 분석
| Dataset | Teacher Mapping (%) | Student Mapping (%) |
|---|---|---|
| SciTail | 71.03 | 80.22 |
| ConTRoL-nli | 68.86 | 75.48 |
| Anli_r2 | 66.37 | 74.47 |
Table 6 (원논문): MinED로 찾은 1:1 토큰 매핑 비율. Bert-base-uncased(student)와 LLM2Vec-Mistral(teacher) 사이의 매핑 결과.
Teacher 기준 6671%, student 기준 7480%의 토큰이 1:1 매핑된다. 양쪽을 평균하면 약 75%인데, 이 정도면 IRA에 투입할 토큰이 충분하다. Student 쪽 매핑 비율이 더 높은 이유는, BERT의 WordPiece 토크나이저가 Mistral의 SentencePiece보다 더 세밀하게 분절하는 경향이 있어서 student 토큰이 teacher 토큰에 포함되는 경우가 많기 때문으로 추정된다.
매핑되지 않는 25~35%의 토큰은 IRA의 어텐션 증류에서 제외되지만, 저자들이 강조하듯 OTIS 단계에서는 전체 시퀀스의 hidden state를 모두 사용하므로 정보 손실이 보완된다. 이 이중 구조가 EMO의 강점이기도 하다.
6.2 Appendix B: 학습 설정 상세
| 설정 항목 | 값 |
|---|---|
| Epoch | 5 |
| Learning Rate | |
| Batch Size | 4 |
| LR Scheduler | cosine |
| Finetune method | LoRA |
| LoRA rank | 256 |
| LoRA alpha | 32 |
| LoRA dropout | 0.1 |
Table 7 (원논문): LLM2Vec Mistral 7B → Bert-base-uncased 증류의 학습 설정.
LoRA rank가 256으로 상당히 높은 편이다. 증류 과정에서 teacher도 함께 fine-tune하는데, rank가 높아야 teacher의 표현력을 충분히 활용할 수 있기 때문으로 보인다.
| Method | Patent | Imdb | Banking77 | SciTail | ConTRoL-nli | Anli_r2 | STSB | STS12 | SICK |
|---|---|---|---|---|---|---|---|---|---|
| OURS | 0.5 | 0.1 | 0.1 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 |
Table 8 (원논문): 각 데이터셋에서 탐색된 최적 하이퍼파라미터 .
9개 데이터셋 중 7개에서 가 최적이다. 태스크 손실과 증류 손실을 동등하게 가중하는 것이 일반적으로 잘 작동한다는 의미다. Imdb와 Banking77만 인데, 이 두 데이터셋은 학습 데이터가 상대적으로 크고(25,000, 10,000) 텍스트 분류 난이도가 낮은 편이라서, student 자체의 학습 신호(CE loss)보다 증류 신호에 더 의존하는 편이 유리한 것으로 해석된다.
6.3 Appendix C: 계산 비용 분석
| Method | ULD | DSKD | MinED | MultiOT | IRA (2 layers) | OTIS | EMO |
|---|---|---|---|---|---|---|---|
| Time (s/batch) | 0.18 | 0.31 | 0.67 | 0.72 | 0.41 | 0.56 | 0.97 |
Table 10 (원논문): SciTail 데이터셋에서 각 방법의 배치당 학습 시간.
EMO의 배치당 학습 시간은 0.97초로, 가장 빠른 ULD(0.18초)의 약 5.4배다. 무시할 수 없는 오버헤드이긴 하지만, 몇 가지 맥락을 고려해야 한다.
첫째, 이 오버헤드는 학습 시에만 발생한다. 추론 시에는 student 모델만 쓰므로 추가 비용이 없다. 둘째, OTIS 단독(0.56초)이 전체 0.97초 중 상당 부분을 차지한다. OT 계산이 비싸서 마지막 레이어에만 적용한 것이 이 수치에서 정당화된다. IRA(0.41초)는 2개 레이어에 적용했는데도 OTIS보다 가볍다. 셋째, MinED(0.67초)와 MultiOT(0.72초)도 이미 비슷한 수준의 계산 비용을 보이는데 성능은 EMO에 한참 못 미친다. 비용 대비 성능(cost-effectiveness) 관점에서 EMO가 합리적인 선택이다.
7. 강점과 한계
강점
- 관계적 + 표현적 지식의 동시 증류: 기존 CTKD가 logit 정렬에만 집중했다면, EMO는 어텐션 패턴(IRA)과 hidden state(OTIS)를 모두 활용한다. Table 3의 ablation에서 두 구성요소가 각각 기여하고 결합 시 시너지가 있음을 확인할 수 있다.
- 태스크 범용성: Logit에 의존하지 않기 때문에 분류(Patent, Imdb 등)와 회귀(STS) 태스크를 모두 커버한다. Table 2에서 보듯 STS 태스크에서는 비교 가능한 CTKD 방법이 DSKD뿐인데, EMO가 일관되게 우위를 보인다.
- Teacher 모델 범용성: Table 5에서 LLM2Vec 외에 BGE-M3 teacher로도 효과를 확인했다. 특정 모델 아키텍처에 종속되지 않는 설계다.
- 중요도 기반 선택적 증류: 모든 토큰을 균등하게 다루는 대신 어텐션 기반 중요도로 가중치를 부여하여 효율과 효과를 동시에 달성한다. Table 4에서 top- 선택이 성능을 높이는 것이 확인된다.
- CKA의 창의적 활용: 기존에 hidden state 비교에 쓰이던 CKA를 어텐션 행렬에 적용한 것은 방법론적으로 참신하다. CKA의 차원 불변 성질( vs 비교 가능) 덕분에 크로스 토크나이저 환경에 자연스럽게 맞는다.
한계 및 아쉬운 점
- 1:1 매핑 제약: MinED의 1:1 매핑에 의존하기 때문에, 매핑되지 않는 토큰(약 25~35%)은 IRA에서 완전히 배제된다. OTIS가 이를 보완한다고는 하지만, 매핑되지 않은 토큰 중에도 구조적으로 중요한 정보가 있을 수 있다. 소프트 매핑이나 다대다(many-to-many) 매핑 전략이 성능을 더 높일 여지가 있다.
- 정적 중요도 점수: 중요도 점수를 teacher의 마지막 레이어 어텐션에서만 계산하고, 학습 과정에서 동적으로 업데이트하지 않는다. Student가 학습됨에 따라 중요한 토큰의 분포가 바뀔 수 있는데, 이를 반영하지 못한다. 저자들도 이를 인식하고 있어 limitations에서 동적/맥락 인지 방법을 향후 과제로 언급했다.
- 실험 다양성의 부족: Student가 Bert-base-uncased 하나로 고정되어 있다. 다른 규모의 student(예: BERT-large, DistilBERT)나 다른 아키텍처 계열(GPT-2, T5 등)에서의 효과는 확인되지 않았다. 또한 retrieval 태스크(MTEB의 다른 카테고리들)에서의 성능이 빠져 있어, 임베딩 모델로서의 범용성을 완전히 검증했다고 보기 어렵다.
- 계산 오버헤드: 배치당 0.97초로 가장 빠른 ULD(0.18초)의 5배 이상이다. 학습 시에만 발생한다고는 하지만, 대규모 데이터에서의 증류 시 무시 못할 비용이다.
- 하이퍼파라미터 에 대한 분석 부족: IRA를 마지막 개 레이어에 적용한다고 했는데, 의 선택 기준이나 민감도 분석이 부록에도 없다. 실험에서는 를 사용한 것으로 보이지만 (Table 10에서 "IRA for 2 layers"), 이 값을 어떻게 정했는지 설명이 없다.
8. 마치며
EMO는 크로스 토크나이저 지식 증류라는 실용적으로 중요한 문제에 대해, 출력 정렬을 넘어 모델 내부 구조까지 전달하는 방법을 제시했다. CKA를 어텐션 행렬에 적용한 IRA와 중요도 기반 mass를 가진 OTIS라는 두 구성요소가 상보적으로 작동하며, logit 의존성을 제거하여 분류와 회귀 태스크 모두에 적용 가능한 범용 프레임워크를 만들었다.
7B teacher에서 110M student로의 63배 압축에서, 여러 태스크에 걸쳐 기존 CTKD 방법을 일관되게 능가하는 결과를 보여준 것은 고무적이다. 특히 STS12에서 baseline 대비 25.6%p, DSKD 대비 10.0%p 향상은 이 프레임워크의 효과를 단적으로 보여준다.
향후에는 1:1 매핑의 제약을 풀고, 중요도 점수를 동적으로 업데이트하며, 더 다양한 모델 조합과 태스크에서 검증하는 방향이 자연스러운 후속 연구가 될 것이다. 실용적 관점에서는 RAG 파이프라인이나 실시간 검색 시스템에서 대형 임베딩 모델을 소형으로 교체할 때, EMO 같은 방법이 품질 손실을 최소화하는 열쇠가 될 수 있다.
References
- Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the knowledge in a neural network.
- Jiao, X., et al. (2019). TinyBERT: Distilling BERT for natural language understanding.
- Wang, W., et al. (2020). MiniLM: Deep self-attention distillation for task-agnostic compression.
- Kornblith, S., et al. (2019). Similarity of neural network representations revisited.
- Boizard, N., et al. (2025). Towards cross-tokenizer distillation: The universal logit distillation loss for LLMs. (ULD)
- Wan, F., et al. (2024). Knowledge fusion of large language models. (MinED)
- Zhang, S., et al. (2024). Dual-space knowledge distillation for large language models. (DSKD)
- Cui, X., et al. (2024). Multi-level optimal transport for universal cross-tokenizer knowledge distillation. (MultilevelOT)
- BehnamGhader, P., et al. (2024). LLM2Vec: Large language models are secretly powerful text encoders.
- Chen, J., et al. (2024). BGE M3-embedding: Multi-lingual, multi-functionality, multi-granularity text embeddings.
- Cuturi, M. (2013). Sinkhorn distances: Lightspeed computation of optimal transportation distances.
- Lin, Z., et al. (2025). Rho-1: Not all tokens are what you need.