kangnlp

논문 리뷰

논문 리뷰: TALAS - Teacher-Anchored Layer Alignment with Adaptive Sharpness-Aware Minimization for Embedding Distillation

36분 읽기

캐싱해둔 교사 임베딩만으로 LLM 임베딩 모델을 BERT급 소형 인코더로 증류하되, 상위 레이어만 교사에 앵커링하고 나머지는 위→아래로 자기 증류하며, ASAM으로 평평한 최소점을 찾아 일반화를 끌어올린 프레임워크.

논문 정보

항목내용
제목TALAS: Teacher-Anchored Layer Alignment with Adaptive Sharpness-Aware Minimization for Embedding Distillation
저자Quoc Phong Dao*, Hoang Son Nguyen*, Pham Khanh Chi*, Linh Ngo Van† (Hanoi University of Science and Technology), Diep Thi-Ngoc Nguyen (VNU-UET), Thien Huu Nguyen (University of Oregon), Trung Le (Monash University)
학회/저널arXiv preprint, 2026년 6월 (cs.CL)
논문 링크arXiv:2606.21851

( 공동 1저자, † 교신저자)*


1. 들어가며

RAG 파이프라인이 실무의 기본기가 되면서, 검색 품질을 좌우하는 텍스트 임베딩 모델의 중요성은 두말할 필요가 없어졌다. 문제는 요즘 성능 상위권을 휩쓰는 임베딩 모델들이 하나같이 파라미터가 크고 임베딩 차원도 높다는 데 있다. MTEB 리더보드 상단을 보면 수십억 파라미터짜리 LLM 기반 임베딩 모델이 즐비한데, 이런 모델을 실제 서비스에 그대로 얹으면 지연시간과 서빙 비용이 감당하기 어려운 수준으로 뛴다. 그래서 "성능은 최대한 가져오되 크기는 줄이자"는 지식 증류(Knowledge Distillation, KD)가 자연스럽게 후보로 떠오른다.

그런데 KD 연구의 무게중심은 그동안 생성형 LLM 쪽에 쏠려 있었고, 정작 다운스트림 곳곳에서 쓰이는 임베딩 모델의 증류는 상대적으로 덜 파고들어졌다. 임베딩 증류에는 고유한 골칫거리가 있기 때문이다. 하나는 효율성과 효과성의 트레이드오프다. 교사의 hidden state나 attention 패턴 같은 내부 신호를 끌어다 쓰면 학생이 배울 감독 신호가 풍부해지지만, 교사가 수십억 파라미터짜리 LLM이면 그 내부 신호를 학습 중에 계속 뽑아내는 비용이 어마어마하다. 반대로 교사의 최종 출력 임베딩 하나만 쓰면 파이프라인은 가벼워지지만, 그 한 벡터에 눌러 담긴 의미를 학생이 온전히 복원하기가 쉽지 않다.

또 하나는 용량 격차(capacity gap) 문제다. 교사가 지나치게 강하면 오히려 증류가 잘 안 되는 역설이 잘 알려져 있다(Cho and Hariharan, 2019). 특히 임베딩 모델에서 학생의 얕은 레이어까지 교사의 고도로 추상화된 표현에 억지로 맞추려 하면, 저수준 언어 특징을 뽑아야 할 초기 레이어의 역할이 망가지면서 최적화가 오히려 불안정해진다.

TALAS는 이 두 문제를 정면으로 겨냥한다. 교사 임베딩을 학습 전에 한 번만 뽑아 캐싱해두고(학습 루프에서 교사 추론은 아예 없다), 그 임베딩을 학생의 상위 레이어에만 앵커로 걸며, 나머지 레이어는 학생 자신의 상위 레이어를 길잡이 삼아 위에서 아래로 구조를 전파한다. 마지막으로 Adaptive Sharpness-Aware Minimization(ASAM)으로 손실 지형의 평평한 최소점을 찾아 일반화를 안정화한다. 세 축이 맞물리면서, 4B짜리 Qwen3 임베딩 모델을 109M BERT-base로, 혹은 0.6B 모델을 22M MiniLM으로 증류하면서도 강력한 베이스라인을 일관되게 앞서는 결과를 낸다.


2. 기존 연구의 한계

논문의 Related Work(본문 2절과 Appendix A)는 증류 방식을 세 갈래로 정리한 뒤, 임베딩 특화 증류를 별도로 다룬다.

2.1 출력 증류 (Output Distillation)

가장 널리 쓰이는 형태로, 학생이 교사의 출력 분포나 출력 표현을 그대로 흉내 내도록 학습한다(Hinton et al., 2015; Sanh et al., 2020). 교사 내부 상태에 접근할 필요가 없어 오프라인·블랙박스 증류에 자연스럽게 들어맞고, 구현이 단순하며 효율적이라는 게 장점이다. 문제는 용량 격차가 클 때다. 출력 하나만 보고 배우면 학생이 교사 특유의 아티팩트에 과적합하거나, 교사 분포의 지배적이고 쉬운 부분만 골라 배우면서 미묘한 뉘앙스를 놓치기 쉽다(Cho and Hariharan, 2019; Gu et al., 2024).

2.2 관계 증류 (Relational Distillation)

출력 하나로는 부족하니, 개별 표현을 맞추는 대신 샘플 사이의 구조적 관계를 전달하자는 흐름이다. Relational Knowledge Distillation(RKD, Park et al., 2019)이 대표적으로, 임베딩 공간에서 쌍별 거리나 각도 관계를 보존한다. 표현 하나하나를 맞추는 것보다 의미 기하(semantic geometry)를 보존하는 편이 더 강건하다는 게 여러 연구에서 확인됐다(Jiao et al., 2020; Truong et al., 2025). 다만 쌍별 비교나 중간 상태 감독이 들어가면 메모리·연산 부담이 커지고, 특히 학습 중 교사 추론을 온라인으로 돌려야 하면 확장성이 크게 떨어진다.

2.3 교사 보조자 증류 (Teacher Assistant Distillation)

용량 격차 자체를 완화하려고 교사와 학생 사이에 중간 용량의 보조 모델을 끼워 넣는 접근이다(Mirzadeh et al., 2019). 표현 차이를 단계적으로 줄여주니 교사가 학생보다 훨씬 클 때 직접 증류보다 일관되게 낫다. 대신 보조 모델을 따로 학습·유지해야 하므로 학습 복잡도와 비용이 늘어난다는 게 근본적 제약이다.

2.4 임베딩 특화 증류

임베딩 교사는 대개 학생과 아키텍처·토크나이저·어휘가 다른 대형 LLM이라, 토큰 단위 정렬 자체가 비싸거나 불가능한 경우가 많다. 그래서 임베딩 전용 증류 연구는 상대적으로 드물다. DistillCSE(Xu et al., 2023)는 대조 학습으로 문장 표현을 전달해 토큰 단위 감독을 피했고, Jasper & Stella(Zhang et al., 2025a)는 다단계 증류 파이프라인으로 SOTA 임베딩 모델을 압축했으며, EMO(Truong et al., 2025)는 최적 수송(optimal transport)으로 샘플 간 관계를 정렬한다.

이들 대부분이 여전히 토큰 단위 감독, 중간 상태 매칭, 반복적인 온라인 교사 추론 중 하나에 의존한다는 게 공통된 한계다. TALAS의 출발점이 바로 여기다. "교사의 캐싱된 문장 임베딩만 쓰면서, 얕은 레이어를 억지로 교사에 맞추지 않고도 용량 격차를 넘을 수 있을까?"


3. 핵심 아이디어

TALAS의 통찰은 "교사를 어디에, 어떻게 걸 것인가"를 다시 설계한 데 있다. 세 가지 결정이 맞물린다.

첫째, 교사는 상위 레이어에만 건다. 저수준 언어 특징을 담당하는 얕은 레이어까지 교사의 추상적 임베딩에 맞추려 하면 오히려 학습이 망가진다. 그래서 학생이 교사의 임베딩 공간을 근사할 만한 용량을 이미 갖춘 상위 레이어에만 교사 앵커링을 제한한다.

둘째, 나머지는 학생 스스로 위→아래로 전파한다. 상위 레이어를 교사에 앵커링했으니, 그 바로 아래 레이어는 교사가 아니라 자기 위층(l+1l+1)을 "길잡이"로 삼아 정렬한다. 이 사슬형 전파가 얕은 레이어를 교사에 직접 노출시키지 않으면서도, 얕은 층부터 깊은 층까지 표현 기하가 매끄럽게 이어지도록 다리를 놓는다. 여기서 점 단위가 아니라 **샘플 간 관계 구조(RKD)**를 인접 레이어끼리 맞춘다는 점이 중요하다. 절댓값 표현이 아니라 "배치 안에서 어떤 샘플이 어떤 샘플과 가깝다"는 위상만 전달하니, 얕은 레이어의 고유한 특징 추출을 파괴하지 않는다.

셋째, 평평한 최소점을 노린다. 코사인 기반 임베딩 증류는 인스턴스 단위의 강한 제약이라, 표준 SGD로 밀어붙이면 학생이 교사 출력에 과적합하면서 데이터의 밑바탕 구조를 못 배우는 위험이 있다(Stanton et al., 2021). ASAM으로 손실 지형의 sharpness를 억제하면, 학생은 교사의 넓은 의미 구조를 붙잡되 특정 임베딩 타깃에 과하게 매몰되지 않는다.

여기에 실무적으로 결정적인 한 가지가 더 붙는다. 교사 임베딩을 학습 전에 딱 한 번 뽑아 캐싱해두므로, 학습 루프 안에서 교사 추론이 완전히 사라진다. 이게 왜 큰가는 4절에서 저장 비용 계산과 함께 다시 짚는다.

기존 임베딩 증류가 "무엇을 얼마나 풍부하게 흉내 낼까"에 집중했다면, TALAS는 "어느 깊이에, 어떤 형태의 신호를, 어떤 최적화로" 전달할지를 분리해서 설계한 것이 근본적 차이다.


4. 제안 방법 (Method)

4.1 전체 구조

Figure 1: An illustration of the proposed TALAS framework.
Figure 1: An illustration of the proposed TALAS framework.

Figure 1: TALAS 프레임워크 개요. 교사(왼쪽)는 배치 입력에서 문장 임베딩을 만들고, 이 임베딩이 학습 가능한 Linear Projection을 거쳐 학생의 상위 레이어(Layer L, L-1)에 LTAMD\mathcal{L}_{TAMD}로 앵커링된다. 학생 내부의 인접 레이어끼리는 LLASD\mathcal{L}_{LASD}로 기하 구조 일관성을 강제한다. 전체 학습은 ASAM으로 최적화되어 평평한 최소점(오른쪽 위)을 향한다. 오른쪽 아래는 LLASD\mathcal{L}_{LASD}(각도 정렬, 관계 보존)와 LTAMD\mathcal{L}_{TAMD}(투영 후 점 단위 정렬)의 개념도. (원논문)

이 그림 한 장에 TALAS의 설계가 압축돼 있다. 왼쪽 파란 블록이 교사인데, 화살표가 배치 입력에서 문장 임베딩으로 한 번만 흐른다는 게 핵심이다 - 학습 중 교사는 다시 호출되지 않는다. 교사 임베딩은 두 개의 서로 다른 Linear Projection(노란 사다리꼴)을 거쳐 학생의 최상위 두 레이어(Layer L, Layer L-1)에 각각 LTAMD\mathcal{L}_{TAMD}로 연결된다. 같은 교사 임베딩을 여러 상위 레이어가 공유 앵커로 삼는 구조다.

학생 내부(보라색 스택)를 보면, 인접 레이어 사이마다 LLASD\mathcal{L}_{LASD} 화살표가 걸려 있다. Layer 1↔2, Layer 2↔3처럼 위층이 아래층을 끌어주는 사슬이다. 맨 아래 "Sentence Embeddings at Layer 1st"까지 이 사슬로 교사의 구조가 전파된다.

오른쪽 위의 손실 지형 그림이 ASAM의 직관을 보여준다. 붉은 SGD 경로는 좁고 깊은 sharp minimum(일반화 나쁨)으로 빨려 들어가는 반면, 초록 ASAM 경로는 넓고 평평한 flat minimum(일반화 좋음)으로 향한다. 오른쪽 아래 두 개의 구(球) 그림은 두 손실의 성격 차이를 대비한다 - LLASD\mathcal{L}_{LASD}는 벡터들의 상대적 각도 배치(관계)를 맞추고, LTAMD\mathcal{L}_{TAMD}는 투영을 통해 개별 벡터를 교사 방향으로 끌어당긴다.

문제 설정을 정리하면 이렇다. 교사 TT와 학생 SS가 있고, 각각의 문장 임베딩을 eTe^T, eSe^S로 쓴다. 학생의 레이어 ll에서 풀링한 임베딩은 elSe^S_l이다. 풀링은 백본에 따라 [CLS](Devlin et al., 2019), [EOS](Zhang et al., 2025b), mean pooling(Lee et al., 2024) 등을 쓴다. 일반적인 KD 손실은 다음처럼 교사·학생 출력의 불일치를 재는 형태다.

LKD=D(zS(x),zT(x))\mathcal{L}_{KD} = D\big(z_S(x),\, z_T(x)\big)

여기서 D(,)D(\cdot,\cdot)는 cross-entropy, MSE, 유사도 기반 거리 등의 불일치 척도다. TALAS는 이 뼈대 위에 세 개의 손실(LTAMD\mathcal{L}_{TAMD}, LLASD\mathcal{L}_{LASD}, LSimCSE\mathcal{L}_{SimCSE})을 얹고 ASAM으로 최적화한다.

4.2 Teacher-Anchored Multi-Layer Distillation (LTAMD\mathcal{L}_{TAMD})

EOFD나 FDD 같은 선행 연구가 attention map, hidden state 같은 복잡한 내부 신호에 기대는 것과 달리, TALAS는 문장 임베딩 증류에만 집중해 연산 부담을 크게 줄인다. 감독 신호가 줄어든 만큼 의미 전달을 극대화하려고, "핵심 의미를 하위 레이어에 압축해 넣자"는 ESE(Li et al., 2025)의 표현 압축 원리에서 아이디어를 빌린다. 그래서 중간·최종 레이어를 모두 교사에 맞추고 싶지만, 용량 격차 탓에 전 레이어를 억지로 맞추는 건 역효과다. 결론적으로 앵커링은 상위 레이어에만 건다.

핵심은 교사의 최종 임베딩을 고정된 앵커로 두고, 학생의 여러 상위 레이어를 그 앵커에 정렬하는 것이다. 레이어 ll에 대한 손실과 전체 손실은 다음과 같다.

LTAMD(l)=1Ni=1NDcos ⁣(el,iSWl,  eiT)\mathcal{L}_{TAMD}(l) = \frac{1}{N}\sum_{i=1}^{N} D_{cos}\!\big(e^S_{l,i}\, W_l,\; e^T_i\big) LTAMD=1k+1l=LkLLTAMD(l)\mathcal{L}_{TAMD} = \frac{1}{k+1}\sum_{l=L-k}^{L} \mathcal{L}_{TAMD}(l)

각 항의 의미는 이렇다. LL은 학생의 전체 레이어 수, NN은 배치 크기, Dcos(u,v)=1uvuvD_{cos}(u,v) = 1 - \frac{u^\top v}{\|u\|\|v\|}는 코사인 거리다. WlRdS×dTW_l \in \mathbb{R}^{d_S \times d_T}는 레이어 ll마다 따로 두는 학습 가능한 투영 행렬로, 학생의 차원 dSd_S를 교사 차원 dTd_T로 매핑한다(그림의 Linear Proj가 바로 이것). 두 번째 식은 최상위부터 k+1k+1개 레이어(l=Lk,,Ll = L-k, \dots, L)의 손실을 평균한다. 즉 kk가 앵커링하는 상위 레이어의 깊이를 정한다.

같은 교사 임베딩 eTe^T가 여러 상위 레이어의 공유 앵커로 작동하는 게 이 설계의 묘미다. 최종 레이어에서만 감독을 주는 대신 여러 상위 레이어를 같은 목표에 걸면, 고수준 의미가 forward pass 도중에 일찍 주입되어 중간 레이어가 미리 의미를 획득한다. 단일 레이어 증류보다 학습이 안정되고 레이어 간 표현 drift가 줄어드는 이유다.

4.3 Layer-Aligned Self-Distillation (LLASD\mathcal{L}_{LASD})

상위 레이어를 교사에 걸었으니 하위 레이어로 감독을 어떻게 이을지가 남는다. 그냥 하위 레이어까지 교사에 직접 맞추면 앞서 말한 대로 얕은 층의 특징 추출이 망가진다(Sun et al., 2019). 그렇다고 하위 레이어를 완전히 방치하면 감독받는 상위 레이어와 방치된 하위 레이어 사이에 표현의 단절이 생겨, 의미 구조가 네트워크를 타고 매끄럽게 흐르지 못한다.

TALAS의 답은 위에서 아래로 순차 전파하는 자기 증류다. 하위 레이어를 교사에 직접 매핑하지 않고, 학생 자신의 위층(l+1l+1)을 아래층(ll)의 동적 길잡이로 쓴다. 이 사슬형 접근 덕에 깊은 층은 고수준 추상을 유지하면서 얕은 층이 자기 구조를 잃지 않고 점진적으로 목표 공간에 다가간다.

이걸 인접 레이어 간의 기하 구조 일관성 문제로 정식화하는데, 여기서 RKD가 등장한다. 점 단위가 아니라 배치 안 샘플들의 구조적 관계를 보존하는 방식이다. 레이어 ll의 배치 임베딩을 ElSRN×dE^S_l \in \mathbb{R}^{N \times d}라 하고, L2 정규화한 E~lS\tilde{E}^S_l로부터 관계 행렬을 만든다.

RlS=E~lS(E~lS)R^S_l = \tilde{E}^S_l \big(\tilde{E}^S_l\big)^\top

RlSRN×NR^S_l \in \mathbb{R}^{N \times N}의 각 원소 (RlS)ij(R^S_l)_{ij}ii번째와 jj번째 샘플의 코사인 유사도, 즉 cos(ei,ej)\cos(e_i, e_j)다. 이 행렬이 레이어 ll 임베딩 공간의 기하 위상(topology)을 담는다. 목표는 인접 레이어의 구조적 발산을 최소화하는 것이고, 관계 행렬 차이의 Frobenius norm으로 손실을 정의한다.

LLASD=1L1l=1L11N2Rl+1SRlSF2\mathcal{L}_{LASD} = \frac{1}{L-1}\sum_{l=1}^{L-1} \frac{1}{N^2} \big\| R^S_{l+1} - R^S_l \big\|^2_F

AF2=i,jAij2\|A\|^2_F = \sum_{i,j} A_{ij}^2는 제곱 Frobenius norm으로, 행렬 수준의 MSE 대용이다. 직관적으로 보면, 위층에서 "샘플 A와 B가 가깝고 C는 멀다"는 관계 지도가 아래층에서도 최대한 유지되도록 강제하는 것이다. 절댓값 임베딩을 복사하는 게 아니라 관계 지도만 물려주므로, 각 레이어는 자기 표현의 절대적 스케일과 저수준 특징을 지키면서도 위상만 정렬한다. 이 사슬을 통해 최상단에 걸린 교사 지식이 입력 레이어까지 구조적으로 흘러내리며 용량 격차를 메운다.

4.4 학습 목표와 대조 정규화 (LSimCSE\mathcal{L}_{SimCSE})

LTAMD\mathcal{L}_{TAMD}LLASD\mathcal{L}_{LASD}는 모두 학생을 교사의 의미 공간에 맞추는 데 초점이 있다. 그런데 직접 감독에만 전적으로 기대면 위험이 있다. 교사 표현이 해당 학습 도메인에서 붕괴(collapse)하거나 이방성(anisotropy)이 심하면, 학생이 그 나쁜 기하 성질까지 물려받아 표현 공간이 쪼그라들 수 있다(Ethayarajh, 2019). 대조 학습은 uniformity를 최적화해 이 붕괴를 막는데, 서로 다른 샘플을 초구면 위에 고르게 흩뿌려 표현이 이방성 클러스터로 퇴화하는 걸 막는다(Wang and Isola, 2020).

그래서 SimCSE(Gao et al., 2022) 기반 비지도 대조 손실을 정규화 항으로 넣는다. 배치 문장을 서로 다른 dropout 마스크 z,zz, z'로 두 번 통과시켜 두 뷰 eize^z_i, eize^{z'}_i를 얻고 다음을 최소화한다.

LSimCSE=1Ni=1Nlogesim(eiz,eiz)/τj=1Nesim(eiz,ejz)/τ\mathcal{L}_{SimCSE} = -\frac{1}{N}\sum_{i=1}^{N} \log \frac{e^{\,\text{sim}(e^z_i,\, e^{z'}_i)/\tau}}{\sum_{j=1}^{N} e^{\,\text{sim}(e^z_i,\, e^{z'}_j)/\tau}}

NN은 배치 크기, τ\tau는 온도, sim()\text{sim}(\cdot)은 코사인 유사도다. 이 항이 교사의 의미 안내를 받으면서도 넓고 균일한 표현 공간을 유지하도록 균형을 잡아준다.

최종 목표는 세 손실의 가중합이다.

Ltotal=λ1LSimCSE+λ2LTAMD+λ3LLASD\mathcal{L}_{total} = \lambda_1 \mathcal{L}_{SimCSE} + \lambda_2 \mathcal{L}_{TAMD} + \lambda_3 \mathcal{L}_{LASD}

λ1,λ2,λ3\lambda_1, \lambda_2, \lambda_3는 세 항의 기여를 조절하는 하이퍼파라미터로, 실험에서는 λ1=0.001\lambda_1 = 0.001, λ2=0.75\lambda_2 = 0.75, λ3=1\lambda_3 = 1을 썼다(Appendix E). 대조 항의 가중치가 유독 작은 게 눈에 띄는데, 이 항은 어디까지나 표현 붕괴를 막는 안전장치 역할이고 실제 의미 전달은 두 증류 항이 주도한다는 설계 의도가 숫자에 드러난다.

4.5 ASAM 최적화

목표 함수가 학생을 교사에 잘 정렬시켜도, 표준 SGD로 증류 손실을 밀면 위험하다. 코사인 기반 임베딩 증류는 인스턴스 단위의 강한 제약이라 학생이 데이터 구조를 못 배운 채 교사 출력에만 과적합할 수 있고, 이는 교사에 대한 충실도는 높지만 일반화는 나쁜 결과로 이어진다(Stanton et al., 2021).

그래서 Ltotal\mathcal{L}_{total}을 ASAM(Kwon et al., 2021)으로 최적화한다. 경험적 손실 L(w)\mathcal{L}(w)를 직접 줄이는 대신 국소 min-max 문제를 푼다.

minw  maxTw1ϵ2ρ  L(w+ϵ)\min_{w}\; \max_{\|T_w^{-1}\epsilon\|_2 \le \rho}\; \mathcal{L}(w + \epsilon)

ww는 모델 파라미터, ϵ\epsilon은 적대적 섭동(perturbation), ρ\rho는 이웃 크기를 조절하는 반경, TwT_w는 파라미터 크기에 따라 섭동을 정규화하는 파라미터 의존 스케일 행렬이다. 바로 이 TwT_w가 SAM과 ASAM을 가르는 지점이다. 표준 SAM이 유클리드 이웃에서 섭동을 재는 반면, ASAM은 파라미터 스케일을 반영해 섭동을 정규화하므로 더 안정적이고 적응적이다. 강한 감독 신호나 큰 용량 격차가 있는 증류에서 특히 유용한 성질이다.

실제 최적화는 2단계 상승-하강 절차로 이뤄진다. Algorithm 1을 풀어 쓰면 이렇다.

  • 사전 단계: 학습 시작 전, 교사 임베딩을 전 학습 데이터에 대해 한 번만 뽑아 캐시 DeT\mathcal{D}^T_e에 저장한다.
  • 각 스텝마다:
    1. 미니배치 BB를 뽑고, 캐시에서 대응하는 교사 타깃 ETE^T를 꺼낸다.
    2. 상승(Ascent): 현재 가중치 ww에서 Ltotal\mathcal{L}_{total}의 그래디언트로 섭동된 가중치 w~\tilde{w}를 만든다(w~ASAM_Perturb(w,wL,ρ)\tilde{w} \leftarrow \text{ASAM\_Perturb}(w, \nabla_w \mathcal{L}, \rho)).
    3. 하강(Descent): 섭동 위치 w~\tilde{w}에서 그래디언트 gASAM=w~Ltotalg_{ASAM} = \nabla_{\tilde{w}} \mathcal{L}_{total}를 다시 계산해 원래 파라미터 ww를 업데이트한다.

요약하면, "지금 위치가 조금만 흔들어도 손실이 확 튀는 뾰족한 골짜기인가?"를 매 스텝 점검하고, 그런 곳을 피해 넓은 분지로 학생을 인도하는 것이다. 대신 스텝마다 forward-backward를 두 번 돌아야 해서 학습 시간이 사실상 두 배가 되는데, 이 비용 얘기는 한계에서 다시 짚는다.

4.6 오프라인 교사 추론과 캐싱된 타깃 (Appendix C)

TALAS의 실무적 무기는 교사를 학습 루프에서 완전히 빼는 데 있다. 교사는 전처리 단계에서 딱 한 번 실행되어 전 학습 샘플의 문장 임베딩(마지막 CLS 임베딩)을 뽑아 저장하고, 학습 중에는 미니배치에 해당하는 사전 계산 임베딩만 lazy loading으로 불러온다.

왜 이게 큰 차이를 만드는지는 저장 비용 계산이 잘 보여준다. BERT 계열 교사가 샘플 하나당 [256,768][256, 768] 모양의 hidden state를 낸다고 하면 256×768=196,608256 \times 768 = 196{,}608개 원소, float32로 샘플당 약 768KB다. 10만 샘플이면 76GB를 넘어 대부분의 학습 환경에서 비현실적이다. 중간 레이어 hidden state까지 캐싱하는 방식이라면 이 비용이 더 폭증한다. 반면 문장 임베딩만 캐싱하면 샘플당 [768][768] 벡터 하나, float32로 약 3KB다. 10만 샘플이면 약 300MB로 줄어 오프라인 증류가 실용적이고 확장 가능해진다. 이 250배 차이가 TALAS가 문장 임베딩만 쓰기로 한 이유를 정량적으로 정당화한다.


5. 실험 결과

5.1 실험 설정

태스크와 데이터셋. 세 부류를 평가한다. Classification은 BANKING77, EMOTION, TWEETEVAL-SENTIMENT에서 F1, Pair Classification은 MRPC, SCITAIL, WIC에서 Average Precision(AP), STS는 SICK, STS12, STS-BENCHMARK에서 Spearman 상관을 잰다. 총 9개 데이터셋이다. EMO(Truong et al., 2025)의 태스크 선정 프로토콜을 따르되, EMO가 지도 학습(SFT)을 쓴 것과 달리 TALAS는 라벨 없이 SimCSE 스타일 비지도 목표로 학습한다.

학습 코퍼스. 세 in-domain 데이터셋(EMOTION, WIC, STS-BENCHMARK)에서 각 5,000문장씩 균등 샘플링해 하나의 라벨 없는 코퍼스로 합친다. 나머지 데이터셋은 out-of-domain으로 취급한다. 즉 학습은 딱 세 데이터셋의 텍스트만 보고, 평가는 9개 전부에서 이뤄지므로 도메인 이동(domain shift)에 대한 강건성을 정직하게 측정하는 구조다. (Appendix D를 보면 실제 코퍼스는 classification 5,000문장 + STS 2,500쌍 + pair classification 2,500쌍을 개별 문장으로 펼쳐 총 15,000문장으로 구성한다.)

모델. 학생은 BERT-base, MINILMV2 H384, MINILMV2 H768. 교사는 BGE-M3, QWEN3-EMBEDDING-0.6B, QWEN3-EMBEDDING-4B. 세 가지 교사→학생 조합을 체계적으로 평가한다: Qwen3-0.6B→MiniLMv2 H384(22M), BGE-M3→MiniLMv2 H768(66M), Qwen3-4B→BERT-base(109M).

베이스라인. 토큰 단위 KD로 DSKD, CDM, EMO를, 문장 단위 KD로 JASPER & STELLA, DISTILLCSE를 비교한다. 여기에 증류 없이 SimCSE만 돌린 SimCSE-unsup과, 학습 전 학생(Student base)도 참조점으로 둔다.

5.2 주요 결과 (Main Results)

Table 1은 세 교사-학생 조합 각각에서 9개 데이터셋 성능과 in/out-of-domain 평균을 보여준다. in-domain은 ⋆로 표시(Emotion, WiC, STSB)했다.

Qwen3-Embedding 0.6B → MiniLMv2 H384 (22M)

MethodBanking77TweetEmotion⋆MRPCSciTailWiC⋆SICKSTS12STSB⋆Avg-InAvg-OutAvg.
Teacher93.2171.3366.5383.3790.0066.8580.6477.1184.6072.6682.6179.29
Student base68.2941.0769.9078.3964.5659.5847.6021.9522.0840.9158.4552.60
SimCSE-unsup86.6769.5654.0082.2973.7665.7166.0459.1162.0860.6072.9168.80
CDM86.0770.3253.0481.9772.9768.3465.9761.1265.7562.3873.0769.51
DSKD85.6669.8652.2582.1773.5268.5866.5462.5266.7062.5173.3869.76
Jasper & Stella85.9471.2557.9983.3676.1768.2570.5763.8170.1165.4575.1871.94
DistillCSE88.0369.8754.2483.2077.6667.8870.0968.1971.5364.5576.1772.30
EMO87.0371.6759.6383.4578.1768.8970.8165.7871.4266.6576.1572.98
TALAS86.6972.7760.9485.1081.4266.6572.5170.7776.2267.9478.2174.79

BGE-M3 → MiniLMv2 H768 (66M)

MethodBanking77TweetEmotion⋆MRPCSciTailWiC⋆SICKSTS12STSB⋆Avg-InAvg-OutAvg.
Teacher93.5273.8568.5685.8191.8761.4779.1878.7384.8771.6383.8379.76
Student base81.8447.7071.6779.6765.0660.8949.4626.7624.1244.2462.4156.35
SimCSE-unsup87.8971.4555.6583.8776.3267.9666.9359.6663.9762.5374.3570.41
CDM89.1570.4858.7184.1675.9170.0068.1462.7467.7865.5075.1071.90
DSKD89.5170.5357.2784.3975.3669.9668.8166.7870.0565.7675.9072.52
Jasper & Stella88.3874.1362.8185.9880.3368.3874.3566.8074.5568.5878.3375.08
DistillCSE90.5072.7059.8884.3978.4869.4973.9770.0073.7367.7078.3474.79
EMO90.1974.2661.4885.3180.8268.3075.4167.0676.3568.7178.8475.46
TALAS90.6274.8163.3087.0581.9466.2876.8669.0479.0169.5380.0576.55

Qwen3-Embedding 4B → BERT-base (109M)

MethodBanking77TweetEmotion⋆MRPCSciTailWiC⋆SICKSTS12STSB⋆Avg-InAvg-OutAvg.
Teacher93.6372.2268.9484.7991.7766.8083.4382.5586.8774.2084.7381.22
Student base84.8647.7968.0277.3667.7459.2242.4321.5420.3042.4460.3354.36
SimCSE-unsup89.2169.1353.5582.8676.5971.6468.1961.6870.3665.1874.6171.47
CDM89.7570.5853.1184.8375.1971.0669.3768.4873.9266.0376.3772.92
DSKD89.6669.9354.5183.1377.9571.5368.6563.3771.4265.8275.4572.24
Jasper & Stella89.7173.8365.9885.3380.9970.5073.9667.9075.6570.7178.6275.98
DistillCSE90.9470.9060.8082.8678.9868.6375.1272.6177.0868.8478.5775.32
EMO90.7873.5467.4884.4681.5769.7675.6668.8377.1771.4779.1476.58
TALAS91.4374.3070.7186.4782.6669.2478.3875.4080.8873.6181.4478.83

Table 1 (원논문): 9개 데이터셋 실험 결과. ⋆는 in-domain, 나머지는 out-of-domain. Avg.는 전체 평균이며 TALAS를 볼드 처리.

세 조합 모두에서 TALAS가 전체 평균(Avg.) 1위를 가져간다. 22M 조합에서 74.79로 2위 EMO(72.98)를 1.81점, 66M에서 76.55로 EMO(75.46)를 1.09점, 109M에서 78.83으로 EMO(76.58)를 2.25점 앞선다. 강한 베이스라인인 EMO가 SFT(지도 학습)를 쓰는 반면 TALAS는 라벨 없는 비지도 학습이라는 점을 감안하면, 이 격차의 의미가 더 커진다.

특히 눈에 띄는 건 out-of-domain(Avg-Out) 성능이다. 109M 조합에서 TALAS의 Avg-Out은 81.44로, 교사(84.73)에 3.3점 차까지 따라붙는다. STS12를 보면 TALAS(75.40)가 EMO(68.83)를 6.6점, Tweet에서도 대부분 베이스라인을 앞선다. 학습 코퍼스가 세 in-domain 데이터셋뿐인데 out-of-domain에서 이만큼 벌린다는 건, 교사 출력을 그대로 복사하는 게 아니라 표현의 기하 구조를 배우는 LLASD\mathcal{L}_{LASD}와 평평한 최소점을 찾는 ASAM이 실제로 일반화에 기여하고 있다는 방증이다.

한 가지 흥미로운 예외는 WiC⋆다. 세 조합 모두에서 TALAS의 WiC 점수가 여러 베이스라인보다 오히려 낮다(22M에서 66.65 vs DSKD 68.58). WiC는 문맥에 따라 같은 단어의 의미가 바뀌는지를 판단하는 태스크라, 문장 전체 임베딩의 의미 유사도만으로는 잡기 어려운 세밀한 어휘 중의성 문제다. 문장 임베딩만 증류하는 TALAS의 설계가 이 지점에서는 상대적으로 불리하게 작용하는 것으로 보인다. 그래도 나머지 8개 데이터셋에서의 우위가 이를 상쇄하고도 남는다.

5.3 Ablation Study - 손실 조합의 기여

Table 2는 네 요소(LSimCSE\mathcal{L}_{SimCSE}, LTAMD\mathcal{L}_{TAMD}, LLASD\mathcal{L}_{LASD}, ASAM)를 켜고 끄며 각각의 기여를 분해한다.

Qwen3-Embedding 0.6B → MiniLMv2 H384 (22M)

LSimCSE\mathcal{L}_{SimCSE}LTAMD\mathcal{L}_{TAMD}LLASD\mathcal{L}_{LASD}ASAMAvg-InAvg-OutAvg-All
44.2462.4156.35
65.5075.2471.99
67.0375.7772.85
67.6577.1073.95
67.2777.7874.28
68.0476.9373.97
67.9478.2174.79

BGE-M3 → MiniLMv2 H768 (66M)

LSimCSE\mathcal{L}_{SimCSE}LTAMD\mathcal{L}_{TAMD}LLASD\mathcal{L}_{LASD}ASAMAvg-InAvg-OutAvg-All
40.9158.4552.60
70.0879.2176.17
69.3779.6376.21
69.4679.1275.90
68.6779.0675.59
69.8379.2776.13
69.5380.0976.57

Qwen3-Embedding 4B → BERT-base (109M)

LSimCSE\mathcal{L}_{SimCSE}LTAMD\mathcal{L}_{TAMD}LLASD\mathcal{L}_{LASD}ASAMAvg-InAvg-OutAvg-All
65.1874.6171.47
69.5877.5374.88
71.2877.9975.75
73.0580.2977.88
72.3280.9778.09
73.3280.4478.06
73.6181.4478.83

Table 2 (원논문): 손실 조합별 ablation. in-domain, out-of-domain, 전체 평균으로 보고. 마지막 행이 전체 TALAS.

이 표가 각 요소의 역할을 또렷이 드러낸다. 먼저 LTAMD\mathcal{L}_{TAMD}가 절대적 뼈대다. 첫 행(SimCSE만)은 사실상 증류 없는 상태로, 22M에서 56.35에 그친다. LTAMD\mathcal{L}_{TAMD}를 켜는 순간 71.99로 15점 이상 뛴다. 논문 표현대로 LTAMD\mathcal{L}_{TAMD}가 빠진 구성은 "파국적 성능 하락"을 겪으며, 다른 목표들이 그 위에서 작동하는 기반임을 확인시킨다.

ASAM의 기여는 조합마다 다르다. BGE-M3→MiniLMv2 H768(66M)에서는 이득이 미미하지만(예: 4행 vs 6행 비교), Qwen3-4B→BERT-base와 Qwen3-0.6B→MiniLMv2 H384에서는 확실히 크다. 109M 조합에서 LSimCSE+LTAMD\mathcal{L}_{SimCSE}+\mathcal{L}_{TAMD}만 쓴 74.88이 ASAM을 더하니(6행) 78.06까지 오른다. 저자의 해석이 설득력 있다 - 학생이 극단적으로 작거나(22M) 교사가 압도적으로 강할 때(4B) 최적화 지형이 험난해지고, 바로 그때 평평한 최소점을 찾는 ASAM의 능력이 결정적이 된다. ASAM은 용량 격차가 클수록 빛을 발하는 장치인 셈이다.

보조 목표 두 개는 역할이 상보적이다. LSimCSE\mathcal{L}_{SimCSE}는 in-domain(Avg-In)의 국소 클러스터링을 강화하지만 학습 분포로 편향되는 조짐이 있고, 대신 LLASD\mathcal{L}_{LASD}를 넣으면 out-of-domain(Avg-Out)이 일관되게 오른다. 109M 조합에서 5행(LTAMD+LLASD+\mathcal{L}_{TAMD}+\mathcal{L}_{LASD}+ASAM, SimCSE 없음)의 Avg-Out은 80.97로, SimCSE를 넣은 6행(80.44)보다 오히려 높다. 레이어 간 기하 관계를 보존하는 게 강건한 전이 학습에 더 효과적이라는 뜻이다. 네 요소를 모두 켠 마지막 행이 세 조합 모두에서 최고 Avg-All을 찍으며, 각 요소가 독립적으로 의미 있는 기여를 한다는 점을 마무리한다.

5.4 계산 효율 분석

Table 3은 Qwen3-4B→BERT-base 조합을 NVIDIA T4 GPU, 배치 32에서 측정한 자원 소비다.

MethodTime (ms/step)Mean Mem (MB)Peak Mem (MB)Std Mem (MB)
SimCSE-unsup15126623690235
DistillCSE18233733674124
Jasper & Stella2042932388471
EMO12231120313831492
DSKD10021099112393336
CDM10781081412156314
TALAS (w/o ASAM)19528643982261
TALAS (with ASAM)37533174499277

Table 3 (원논문): Qwen3-Embedding-4B → BERT-base 학습 시 런타임과 GPU 메모리 비교.

수치가 극적이다. 토큰 단위 베이스라인(EMO, DSKD, CDM)은 학습 중 교사를 병렬로 추론하고 dense state를 캐싱하기 어려운 탓에 스텝당 1,000ms를 넘고 peak 메모리가 12GB를 웃돈다. 반면 TALAS(w/o ASAM)는 스텝당 195ms에 peak 3.98GB로, 토큰 단위 방법 대비 시간은 5배 이상, 메모리는 3배 이상 가볍다. DistillCSE(182ms)에 근소하게 뒤지고 Jasper(204ms)보다는 빠른 수준이면서, Jasper처럼 번거로운 2단계 파이프라인도 필요 없다.

ASAM을 켜면 forward-backward가 두 번이라 375ms로 늘고 peak 4.5GB가 되지만, 그래도 토큰 단위 방법과는 비교 자체가 안 될 만큼 자원 친화적이다. 4.6절의 캐싱 전략(76GB→300MB)이 이 효율의 근본 원천이다. 소비자용 GPU 한 장으로도 4B 교사에서 증류를 돌릴 수 있다는 게 실무적으로 가장 크게 와닿는 지점이다.

5.5 Sharpness-Aware 변형 비교

Figure 2: Comparison of SAM variants (SAM, DISAM, and ASAM).
Figure 2: Comparison of SAM variants (SAM, DISAM, and ASAM).

Figure 2: Qwen3-Embedding 0.6B → MiniLMv2 H384에서 SAM 변형 비교. 파랑=Avg In-Domain, 주황=Avg Out-of-Domain, 초록=Avg All. (원논문)

세 sharpness-aware 최적화(SAM, ASAM, DISAM)를 22M 조합에서 비교한 막대그래프다. ASAM이 전체 평균 74.79로 가장 높고, DISAM(74.37), SAM(73.90)이 뒤를 잇는다. ASAM은 out-of-domain 78.21과 in-domain 67.94를 동시에 최고로 찍어, 일반화와 표현 충실도 사이의 균형이 가장 좋다. DISAM은 도메인별 통계를 섭동 추정에 넣는 더 정교한 변형인데도 ASAM에 근소하게 밀리는데, 이 태스크 구성에서는 파라미터 스케일 정규화(ASAM)만으로도 도메인 간 균형을 충분히 잡는다는 뜻으로 읽힌다. 결국 TALAS는 ASAM을 기본 최적화로 채택한다.

5.6 경험적 Sharpness 분석

sharpness가 실제로 억제되는지를 직접 확인하려고, 학습 중 손실의 곡률을 헤시안 최대 고윳값 λmax\lambda_{max}로 측정했다. 값이 클수록 뾰족한 최소점, 작을수록 평평한 최소점을 뜻하는 널리 쓰이는 프록시다.

EpochASAMAdamW
114.8930.24
312.9421.48
511.5919.66
1010.8518.66
309.9317.41

Table 4 (원논문): 학습 에폭별 헤시안 최대 고윳값 λmax\lambda_{max}. 값이 낮을수록 평평한 최소점.

전 구간에서 ASAM의 λmax\lambda_{max}가 AdamW보다 뚜렷하게 작고, 그 격차가 학습 내내 유지된다(에폭 1에서 14.89 vs 30.24, 에폭 30에서도 9.93 vs 17.41). 단순히 비슷한 해에 도달하되 분산만 줄인 게 아니라, ASAM이 최적화를 체계적으로 평평한 영역으로 몰아간다는 직접 증거다. sharpness-aware 방법이 파라미터 섭동에 대한 민감도를 명시적으로 벌한다는 이론적 동기와 정확히 들어맞는 결과다. Figure 1 오른쪽 위의 손실 지형 그림이 말로 그린 직관을 이 표가 숫자로 뒷받침하는 셈이다.

5.7 증류 레이어 깊이의 효과

Figure 3: Effect of the number of distilled layers.
Figure 3: Effect of the number of distilled layers.

Figure 3: Qwen3-Embedding 4B → BERT에서 교사 앵커링 레이어 수(KK)에 따른 성능. 파랑=Avg In-Domain, 주황=Avg Out-of-Domain, 초록=Avg All. (원논문)

LTAMD\mathcal{L}_{TAMD}로 교사에 앵커링하는 상위 레이어 수 KK를 바꿔가며 본 결과다. 성능이 선형으로 오르지 않고 K=2K=2에서 정점을 찍는다(Avg All 78.83, Avg-Out 81.44). 그 뒤로는 완만하게 떨어지고 K>4K>4에서 out-of-domain이 특히 지속적으로 하락한다. 저자의 해석은 명료하다 - 교사-학생 용량 격차가 큰 상황에서 너무 많은 내부 레이어에 제약을 걸면 과잉 정규화(over-regularization)가 일어나 학생 고유의 특징 조직을 옥죈다는 것이다. 상위 두 레이어만 앵커링하는 게 고수준 의미 전달과 학생의 표현 유연성 사이 최적점이라는 결론으로, 3.1절에서 앵커링을 "상위 레이어에만" 제한한 설계 결정을 실험으로 정당화한다.


6. 추가 실험 및 부록 (Appendix)

6.1 학습 설정 (Appendix E, Table 5)

세 교사-학생 쌍에 통일된 학습 프로토콜을 적용했다. 손실 가중치는 λ1=0.001\lambda_1 = 0.001, λ2=0.75\lambda_2 = 0.75, λ3=1\lambda_3 = 1이다.

SettingsDSKDCDMEMOStella & JasperDistillCSETALAS
Epoch5552+5 (2단계)55
LR2×10⁻⁵2×10⁻⁵2×10⁻⁵2×10⁻⁵2×10⁻⁵2×10⁻⁵
Batch Size323232323232
LR SchedulerCosineCosineCosineCosineCosineCosine

Table 5 (원논문): 증류·문장 임베딩 방법별 학습 설정.

모든 방법이 동일한 5 에폭, LR 2×10⁻⁵, 배치 32, Cosine 스케줄러로 맞춰져 공정 비교가 담보된다. 유독 Stella & Jasper만 "2+5"의 2단계 파이프라인을 요구하는데, 5.4절에서 지적한 이 방법의 번거로움이 여기에도 드러난다. 평가는 classification에 F1, pair classification에 AP, STS에 Spearman 상관을 쓴다. classification은 임베딩을 얼려 로지스틱 회귀를 얹고, pair classification은 코사인 유사도에 최적 임계값을 적용하는 표준 프로토콜(Conneau and Kiela, 2018)을 따른다.

6.2 사전학습 모델과의 비교 (Appendix F, Table 6)

TALAS가 고작 15K 문장으로 학습했다는 점을, 약 100만 규모로 학습한 기성 모델들과 나란히 놓고 본다.

ModelTraining DataBanking77TweetMRPCSciTailSICKSTS12Avg-Out
Unsup-SimCSE~1M sentences90.7469.8584.8583.0472.2368.4078.19
Sup-SimCSE~1M pairs (NLI)90.7774.1086.7685.4079.5277.7682.39
SBERT-NLI (mean)~1M pairs (NLI)88.8474.0086.9776.6769.7170.1577.72
SBERT-NLI+STS (mean)~1M pairs + 8.6k (STS-B)89.8474.1189.2883.7677.1579.3882.25
SBERT-NLI (CLS)~1M pairs (NLI)89.1175.9087.8077.6372.2169.4178.68
TALAS~15K sentences91.4374.3086.4782.6678.3875.4081.44

Table 6 (원논문): 기성 fine-tuned BERT-base 모델과의 out-of-domain 성능 비교. 최고=볼드, 차선=이탤릭.

핵심은 데이터 효율이다. 경쟁 모델 대부분이 약 100만 규모의 감독 데이터로 학습했는데, TALAS는 15K 문장(약 60~70배 적은 규모)만으로 Avg-Out 81.44를 찍어 Sup-SimCSE(82.39), SBERT-NLI+STS(82.25) 같은 지도 학습 SOTA에 근접한다. Banking77에서는 91.43으로 오히려 모두를 앞서고, SICK에서도 78.38로 최상위권이다. STS 태스크에서 강한 걸 보면, 제한된 데이터에서도 효과적인 의미 정렬을 배웠다는 뜻이다. 교사 감독을 효율적으로 활용해 대규모 주석 코퍼스 없이도 고품질 문장 표현을 얻는다는 게 이 표의 메시지다.

6.3 LTAMD\mathcal{L}_{TAMD} 깊이 상세 Ablation (Table 7)

Figure 3을 데이터셋 단위로 펼친 표다.

Num LayersBanking77TweetEmotion⋆MRPCSciTailWiC⋆SICKSTS12STSB⋆Avg-InAvg-OutAvg.
Qwen3-0.6B → MiniLMv2 H384
186.2872.3360.4684.3380.7366.5772.1568.1475.4367.4977.3374.05
286.6972.7760.9485.1081.4266.6572.5170.7776.2267.9478.2174.79
386.0372.3561.0185.1080.5566.8673.0168.6075.7667.8877.6174.36
685.9571.7260.8885.3379.4767.1271.8567.3374.1367.3876.9473.75
Qwen3-4B → BERT-base
191.1072.8870.7385.3682.3071.2777.4065.4178.0973.3679.0877.17
291.4374.3070.7186.4782.6669.2478.3875.4080.8873.6181.4478.83
491.3573.9970.9486.5582.8069.1177.8075.8680.5873.5481.3978.78
691.2474.0969.5986.4582.4869.2677.2572.9179.7272.8680.7478.11
990.0574.8469.1086.7382.3567.9976.7673.1779.0972.0680.6577.79
1290.2874.6868.4786.7482.9968.1576.4572.4378.9671.8680.6077.68

Table 7 (원논문): LTAMD\mathcal{L}_{TAMD}의 앵커링 레이어 수에 대한 ablation.

두 조합 모두 K=2K=2에서 Avg.가 정점이다(74.79, 78.83). 109M 조합은 2~4 레이어(81.44, 81.39)까지 out-of-domain이 최고 수준이지만 6~12로 깊어질수록 79.72→78.96으로 계속 내려간다. in-domain은 상대적으로 안정적인데 out-of-domain이 먼저 무너진다는 게 중요한 관찰이다. 교사 감독을 학생 깊숙이 밀어 넣으면 학습 도메인에는 잘 맞춰도 일반화 능력이 먼저 희생된다는 뜻으로, 앵커링을 소수 상위 레이어로 제한해야 한다는 설계 원칙을 재확인한다.

6.4 LLASD\mathcal{L}_{LASD} 깊이 상세 Ablation (Table 8)

자기 증류를 적용하는 레이어 수를 바꾼 결과인데, LTAMD\mathcal{L}_{TAMD}와 정반대 경향을 보인다.

Num LayersBanking77TweetEmotion⋆MRPCSciTailWiC⋆SICKSTS12STSB⋆Avg-InAvg-OutAvg.
Qwen3-0.6B → MiniLMv2 H384
287.9372.3061.4983.3980.2768.7672.3964.9073.3567.8776.8673.86
387.3472.5760.6685.2381.3067.4371.7769.6175.6667.9277.9774.62
686.6972.7760.9485.1081.4266.6572.5170.7776.2267.9478.2174.79
Qwen3-4B → BERT-base
291.3673.5668.2184.9781.2970.0677.3370.6078.2872.1879.8577.30
491.5974.5070.9686.3081.4269.6377.7674.0479.4073.3380.9478.40
691.6674.5670.9286.2682.4170.0178.0874.8180.3573.7681.3078.78
990.6374.6969.4186.5282.4068.4878.2774.8980.7972.8981.2378.45
1291.4374.3070.7186.4782.6669.2478.3875.4080.8873.6181.4478.83

Table 8 (원논문): LLASD\mathcal{L}_{LASD}의 자기 증류 레이어 수에 대한 ablation.

여기서는 레이어를 깊게 쓸수록 성능이 오른다. 22M 조합은 6 레이어에서 최고 Avg-Out(78.21), 109M 조합은 2→12로 갈수록 단조 증가해 12 레이어에서 최고(Avg-Out 81.44, Avg. 78.83)를 찍는다. LTAMD\mathcal{L}_{TAMD}가 소수 레이어에서 정점을 찍고 꺾이는 것과 대조적이다.

두 표를 나란히 놓으면 TALAS 설계의 논리가 선명해진다. 교사 앵커링(LTAMD\mathcal{L}_{TAMD})은 얕게, 자기 증류(LLASD\mathcal{L}_{LASD})는 깊게. 교사의 추상적 임베딩을 억지로 깊이 밀어 넣으면 학생이 옥죄이지만, 학생 자신의 위층으로 아래층을 끌어주는 관계 정렬은 네트워크 전체에 깊이 깔아도 표현 강건성을 계속 높인다. 큰 학생일수록 깊은 자기 증류의 이득이 더 크다는 점도 확인된다.

6.5 SAM 변형 데이터셋 상세 (Table 9)

Figure 2를 두 조합·전 데이터셋으로 펼친 표다.

MethodBanking77TweetEmotion⋆MRPCSciTailWiC⋆SICKSTS12STSB⋆Avg-InAvg-OutAvg.
Qwen3-0.6B → MiniLMv2 H384
SAM87.5772.7760.6483.5480.4968.0472.3066.0873.6767.4577.1373.90
ASAM86.6972.7760.9485.1081.4266.6572.5170.7776.2267.9478.2174.79
DISAM86.8473.1160.5284.9981.1165.2072.3069.8075.4367.0578.0374.37
Qwen3-4B → BERT-base
SAM91.1673.9768.3985.1080.8470.3877.4878.7775.6571.4781.2277.97
ASAM91.4374.3070.7186.4782.6669.2478.3875.4080.8873.6181.4478.83
DISAM91.2774.6071.1486.5181.9268.5378.5776.0080.5373.4081.4878.79

Table 9 (원논문): SAM 변형(SAM, DISAM, ASAM) 데이터셋별 비교.

ASAM이 두 조합 모두 최고 Avg.(74.79, 78.83)를 유지한다. DISAM이 Tweet 감성 분석 같은 특정 태스크에서는 근소하게 앞서기도 하고, 109M 조합의 Avg-Out에서는 DISAM(81.48)이 ASAM(81.44)을 0.04점 앞서는 접전이지만, 전체 스펙트럼에서의 강건성과 안정성은 ASAM이 우위다. 특정 태스크의 미세한 우열보다 평균적 안정성을 택해 ASAM을 기본 최적화로 삼은 판단이 이 표에 근거를 둔다.

6.6 확장된 관련 연구와 SAM 계열 배경 (Appendix A, B)

Appendix A는 출력·관계·교사 보조자 증류를 2절보다 상세히 정리하며, 임베딩 특화 증류에서 TALAS가 "캐싱된 문장 임베딩만으로 교사 추론 없이 학습한다"는 위치를 재확인한다. Appendix B는 SAM 계열을 정리한다. 표준 SAM(Foret et al., 2020)은 반경 ρ\rho의 유클리드 이웃에서 최악 손실을 최소화하고(minθmaxϵ2ρL(θ+ϵ)\min_\theta \max_{\|\epsilon\|_2 \le \rho} \mathcal{L}(\theta+\epsilon)), 1차 테일러 근사로 최적 섭동을 ϵ^=ρθL/θL\hat{\epsilon} = \rho\,\nabla_\theta \mathcal{L} / \|\nabla_\theta \mathcal{L}\|로 구한다. DISAM(Zhang et al., 2024b)은 SAM의 i.i.d. 가정이 다중 도메인에서 깨진다는 문제의식에서, 도메인 손실 간 분산을 최소화하는 제약을 섭동 추정에 더한다. 잘 수렴한 도메인에는 섭동을 키우고 덜 수렴한 도메인에는 억눌러 도메인 간 수렴 속도를 동기화하는 것이다. TALAS는 이 계열 중 파라미터 스케일을 반영하는 ASAM을 임베딩 증류 파이프라인에 직접 통합해, 강한 감독과 큰 용량 격차 상황에서 강건성과 out-of-domain 일반화를 함께 끌어올린다.

6.7 베이스라인 구현 세부 (Appendix E)

생성형 증류 방법을 판별적 문장 임베딩 설정에 맞게 어떻게 이식했는지도 밝힌다. CDM은 원래 교사의 예측 엔트로피로 DTW 정렬 비용을 조절하는데, 문장 인코더는 토큰 단위 어휘 분포를 내지 않으므로 비용 함수를 토큰 편집 거리(Levenshtein)로 바꿔 DTW로 최적 정렬 경로를 찾고 교사 hidden state를 학생 길이에 맞춰 평균 집계한 뒤 MSE를 건다. DSKD는 양방향·대칭 과정으로 확장해 문장 수준(CLS를 학습 가능한 투영으로 정렬)과 토큰 수준(Cross-Layer Attention으로 어휘·길이 불일치를 우회) 두 granularity에서 정렬한다. 베이스라인을 그냥 가져다 쓴 게 아니라 임베딩 설정에 맞게 성실히 적응시켰다는 점에서, 비교의 공정성을 높이려는 노력이 읽힌다.


7. 강점과 한계

강점

  • 효율과 성능의 동시 달성: Table 3에서 보듯 TALAS는 토큰 단위 방법 대비 학습 시간 5배, 메모리 3배 이상 이득을 보면서도 Table 1에서 전 조합 평균 1위를 가져간다. 교사를 학습 루프에서 빼는 캐싱 전략(76GB→300MB)이 이 효율의 뿌리다.
  • 일반화에 강함: 라벨 없는 15K 문장, 그것도 세 in-domain 데이터셋만으로 학습하고도 out-of-domain에서 격차를 벌린다. Table 6에서 100만 규모 감독 데이터로 학습한 Sup-SimCSE(82.39)에 15K 문장으로 81.44까지 따라붙는 데이터 효율은 특히 인상적이다.
  • 설계와 실험의 정합성: "LTAMD\mathcal{L}_{TAMD}는 얕게, LLASD\mathcal{L}_{LASD}는 깊게"라는 원칙이 Table 7·8의 상반된 깊이 경향으로 깔끔하게 검증된다. ASAM이 sharpness를 실제로 억제한다는 것도 Table 4의 λmax\lambda_{max} 측정으로 직접 확인했다. 주장마다 대응하는 증거가 있다.
  • 재현성 있는 통제 비교: Table 5처럼 모든 방법을 동일 에폭·LR·배치로 맞추고, 생성형 베이스라인도 임베딩 설정에 성실히 이식했다(Appendix E).

한계 및 아쉬운 점

  • 저자가 밝힌 제약: 계산 예산 탓에 대형 모델이나 대규모 코퍼스로의 확장성은 검증하지 못했고, 평가가 영어 중심이라 다국어·저자원 언어에서의 적응성은 미지수다. 또 ASAM이 스텝당 forward-backward를 두 번 요구해 학습 시간이 사실상 두 배가 된다(Table 3에서 195ms→375ms). 수렴 안정성과 학습 속도의 트레이드오프인데, 빠른 반복이 필요한 상황에서는 부담이 될 수 있다.
  • WiC의 일관된 약세: Table 1의 세 조합 모두에서 TALAS의 WiC 점수가 여러 베이스라인에 뒤진다. 문장 임베딩만 증류하는 설계가 어휘 수준 중의성 판단에는 불리하다는 구조적 한계로 보이는데, 논문이 이 지점을 명시적으로 분석하지는 않는다.
  • 평가 태스크의 폭: classification/pair classification/STS 9개에 그치고, 정작 임베딩 모델의 주 무대인 대규모 검색(retrieval) 벤치마크(MTEB 등)가 빠져 있다. 저자도 결론에서 MTEB 확장을 향후 과제로 남겼는데, RAG 실무 관점에서는 검색 성능이 가장 궁금한 지표라 이 공백이 아쉽다.
  • Table 1과 Table 2의 미세한 불일치: ablation(Table 2)의 "LSimCSE\mathcal{L}_{SimCSE}만" 행 값이 Table 1의 Student base 값과 두 조합에서 서로 뒤바뀐 것처럼 보이는 지점이 있다. 최종 TALAS 행은 두 표가 일치하므로 결론에는 영향이 없지만, 편집 단계의 사소한 표기 혼선으로 읽힌다.
  • λ\lambda 민감도 부재: λ1=0.001\lambda_1=0.001처럼 극단적으로 작은 대조 가중치를 썼는데, 이 값들에 대한 민감도 분석이 없어 하이퍼파라미터의 강건성을 판단하기 어렵다.

8. 마치며

TALAS의 기여를 한 문장으로 압축하면, "교사를 어느 깊이에 어떤 형태로 걸 것인가"를 재설계해 임베딩 증류의 효율-성능 트레이드오프를 넘어선 것이다. 상위 레이어에만 교사를 앵커링하고(LTAMD\mathcal{L}_{TAMD}), 나머지는 학생 자신의 위층으로 관계 구조를 위→아래로 전파하며(LLASD\mathcal{L}_{LASD}), ASAM으로 평평한 최소점을 찾는다는 세 결정이 용량 격차라는 고질적 병목을 정면으로 겨냥한다. 여기에 교사 임베딩을 미리 캐싱해 학습 루프에서 교사 추론을 완전히 제거한 실무적 선택이 더해져, 소비자용 GPU 한 장으로 4B 교사를 증류하는 그림이 현실이 됐다.

특히 라벨 없이 15K 문장만으로 100만 규모 지도 학습 모델에 근접하는 데이터 효율은, 주석 비용이 부담스러운 도메인이나 저자원 환경에서 매력적인 대안이 된다. 교사 출력을 복사하는 대신 표현의 기하 구조를 배우게 하고, sharpness를 억눌러 일반화를 확보한다는 두 축은 임베딩 증류를 넘어 다른 표현 학습 증류에도 옮겨볼 만한 아이디어다.

남은 숙제도 분명하다. 검색 벤치마크에서의 검증, 다국어로의 확장, ASAM의 두 배 학습 시간 완화, WiC 같은 어휘 수준 태스크의 약점 보완이 그것이다. 저자가 예고한 MTEB 확장이 이뤄지면 TALAS의 실용적 가치가 훨씬 또렷해질 것이다. 그래도 "무엇을 흉내 낼까"가 아니라 "어디에, 어떻게, 어떤 최적화로"를 분리해 사고한 이 논문의 관점은, 앞으로의 임베딩 증류 연구에 하나의 유용한 좌표를 남긴다.


References

  • Kwon et al., 2021. "ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks."
  • Foret et al., 2020. "Sharpness-Aware Minimization for Efficiently Improving Generalization." (SAM)
  • Zhang et al., 2024b. "Domain-Inspired Sharpness-Aware Minimization." (DISAM)
  • Park et al., 2019. "Relational Knowledge Distillation." (RKD)
  • Hinton et al., 2015. "Distilling the Knowledge in a Neural Network."
  • Gao et al., 2022. "SimCSE: Simple Contrastive Learning of Sentence Embeddings."
  • Xu et al., 2023. "DistillCSE: Distilled Contrastive Learning for Sentence Embeddings."
  • Zhang et al., 2025a. "Jasper and Stella: Distillation of SOTA Embedding Models."
  • Truong et al., 2025. "EMO: Embedding Distillation via Optimal Transport."
  • Cho and Hariharan, 2019. "On the Efficacy of Knowledge Distillation."
  • Mirzadeh et al., 2019. "Improved Knowledge Distillation via Teacher Assistant."
  • Stanton et al., 2021. "Does Knowledge Distillation Really Work?"
  • Wang and Isola, 2020. "Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere."
  • Chen et al., 2024. "BGE-M3: Multi-Linguality, Multi-Functionality, Multi-Granularity Text Embeddings."
  • Zhang et al., 2025b. "Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Models."
  • Li et al., 2025. "ESE: Espresso Sentence Embeddings."