kangnlp

논문 리뷰

논문 리뷰: Following the Autoregressive Nature of LLM Embeddings via Compression and Alignment

23분 읽기

Following the Autoregressive Nature of LLM Embeddings via Compression and Alignment 리뷰

LLM의 자기회귀적 본성을 거스르지 않으면서 고품질 텍스트 임베딩을 생성하는 방법, AutoRegEmbed

논문 정보

항목내용
제목Following the Autoregressive Nature of LLM Embeddings via Compression and Alignment
저자Jingcheng Deng, Zhongtao Jiang, Liang Pang 외 (ICT, CAS / Kuaishou Technology)
학회/저널arXiv 2502.11401 (2025)
논문 링크arXiv
코드GitHub

1. 들어가며

LLM을 텍스트 임베딩 모델로 쓰겠다는 발상은 이제 꽤 익숙하다. BERT 같은 인코더 모델이 오랫동안 텍스트 임베딩의 주류였지만, GPT 계열의 디코더 전용 LLM이 대규모 사전학습을 통해 축적한 언어 이해 능력은 무시할 수 없는 수준에 이르렀기 때문이다. 문제는 그 능력을 "임베딩"이라는 형태로 끌어내는 과정이 생각만큼 깔끔하지 않다는 데 있다.

디코더 전용 LLM의 히든 스테이트는 본질적으로 다음 토큰의 확률 분포를 예측하기 위해 최적화되어 있다. 마지막 토큰의 히든 스테이트를 텍스트 임베딩으로 쓴다? 그 벡터는 입력 텍스트 전체의 의미를 담고 있는 게 아니라, 다음에 올 토큰 하나의 의미를 담고 있을 가능성이 높다. 여기에 대조 학습(Contrastive Learning)을 바로 적용하면, LLM이 사전학습에서 배운 생성적 능력과 대조 학습이 요구하는 판별적 능력 사이에 근본적인 괴리가 생긴다. 이 간극을 메우려면 대량의 학습 데이터와 막대한 GPU 시간이 필요하고, 그래서 7B 규모의 임베딩 모델을 학습시키려면 수백만 개의 트리플렛과 A100 수천 시간이 드는 게 현실이다.

이 논문은 바로 이 지점을 파고든다. "LLM의 자기회귀적 본성을 따르면서(following) 임베딩을 만들 수는 없을까?" 저자들은 AutoRegEmbed라는 방법을 제안하는데, 기존 대조 학습의 코사인 유사도 기반 최적화를 조건부 확률 분포 기반 최적화로 대체한다. 5만 개의 학습 샘플로 수백만 개를 쓰는 SOTA 모델에 필적하는 성능을 달성했다는 게 이 논문의 주장이다.


2. 기존 연구의 한계

2.1 학습 없는 LLM 임베딩 (LLMs without Fine-Tuning)

LLM을 파인튜닝 없이 바로 임베딩 모델로 쓰려는 시도도 꽤 있었다. Echo Embeddings는 텍스트를 자기회귀 모델에 두 번 넣어서 아키텍처적 한계를 보완하지만, 계산 비용이 두 배로 든다. PromptEOL은 "한 단어로 이 문장을 표현하면"이라는 프롬프트를 붙여서 의미 있는 임베딩을 유도하고, MetaEOL은 이를 8가지 프롬프트 유형으로 확장했다. GenEOL은 LLM으로 문장 변환을 생성한 뒤 임베딩을 집계하는 방식을 쓴다. 이런 방법들은 결국 성능이 부족하거나, 여러 번의 추론이 필요해서 실용성이 떨어진다.

2.2 대조 학습 기반 LLM 임베딩 (LLMs with Fine-Tuning)

주류 접근법은 대조 학습으로 LLM을 파인튜닝하는 것이다. RepLLaMA는 LLaMA의 마지막 토큰 히든 스테이트를 DPR 파이프라인에 넣어 대조 학습을 적용했다. LLM2Vec은 단방향 어텐션의 한계를 인식하고 양방향 어텐션 메커니즘을 도입했으며, NV-Embed는 Latent Attention Layer를 추가로 달았다. bge-en-icl은 LLM의 원래 프레임워크를 유지하면서 인컨텍스트 학습을 활용하는 것이 최적이라고 주장한다.

이 방법들의 공통적인 문제는 언어 모델링과 대조 학습 사이의 근본적 차이를 무시한다는 점이다. LLM은 다음 토큰 확률 분포를 모델링하도록 학습되었는데, 대조 학습은 임베딩 간 코사인 거리를 최적화한다. 이 두 목표 사이의 간극이 학습 비효율의 원인이다. Llama2Vec이 비지도 적응 태스크를 먼저 수행한 뒤 대조 학습을 하는 방식으로 이 문제에 가장 가까이 접근했지만, 여전히 최종 단계에서는 코사인 기반 대조 학습에 의존한다.

2.3 생성적 태스크를 활용한 임베딩 개선

인코더 모델 시절에도 생성적 태스크로 임베딩을 개선하려는 시도는 있었다. coCondenser는 비지도 코퍼스로 [CLS] 토큰의 표현력을 높이는 사전학습 전략을 제안했고, PaSeR와 RetroMAE는 인코더-디코더 아키텍처를 활용했다. 하지만 이들은 대조 감독 신호를 생성적 목표에 통합하지 못했고, 소규모 모델에서만 검증되었다는 한계가 있다.


3. 핵심 아이디어

AutoRegEmbed의 출발점은 두 가지 관찰이다.

첫째, LLM의 마지막 토큰 히든 스테이트는 로컬 시맨틱(다음 토큰의 의미)을 담고 있지, 글로벌 시맨틱(입력 전체의 의미)을 담고 있지 않다. 평균 풀링이나 어텐션 풀링으로 이를 보완하려는 시도가 있지만, 평균 풀링은 "볼록 보존(convexity preservation)" 메커니즘에 불과하고, 어텐션 풀링은 원래 아키텍처를 변형해서 사전학습의 이점을 훼손한다.

둘째, 대조 학습의 코사인 유사도 최적화는 LLM의 사전학습 목표(다음 토큰 확률 분포 모델링)와 근본적으로 다른 최적화 공간에서 작동한다. 이 괴리 때문에 LLM이 사전학습에서 배운 능력을 제대로 활용하지 못한다.

AutoRegEmbed는 이 두 문제를 각각 **정보 압축(Information Compression)**과 **조건부 분포 정렬(Conditional Distribution Alignment)**로 해결한다. 정보 압축은 컨텍스트 압축 기법에서 영감을 받아, 입력 텍스트의 글로벌 시맨틱을 소수의 압축 토큰에 집약한다. 조건부 분포 정렬은 InfoNCE의 구조를 차용하되, 코사인 유사도 대신 조건부 확률 분포 간의 거리를 유사도 메트릭으로 사용한다. 두 태스크 모두 자기회귀적 프레임워크 안에서 작동하기 때문에, LLM의 사전학습 능력을 최대한 보존하면서 임베딩 품질을 높일 수 있다는 것이 저자들의 주장이다.


4. 제안 방법 (Method)

4.1 전체 구조

AutoRegEmbed는 두 단계로 구성된다. 1단계 정보 압축에서 입력 텍스트의 글로벌 시맨틱을 압축 토큰에 집약하고, 2단계 조건부 분포 정렬에서 이 압축 토큰을 임베딩으로 사용하여 양성/음성 샘플과의 분포 정렬을 수행한다.

Figure 2: AutoRegEmbed 전체 프레임워크
Figure 2: AutoRegEmbed 전체 프레임워크

Figure 2: AutoRegEmbed의 전체 프레임워크. 왼쪽은 정보 압축 태스크로 컨텍스트와 인스트럭션을 압축 토큰에 집약하고, 오른쪽은 조건부 분포 정렬 태스크로 양성 샘플과의 분포를 일치시키면서 음성 샘플과는 멀어지게 학습한다. 인코더와 디코더는 동일한 구조를 공유한다. (원논문)

Figure 2는 이 두 단계의 전체 흐름을 보여준다. 왼쪽 패널에서 LLM 인코더는 컨텍스트, 인스트럭션, 압축 토큰을 입력받아 임베딩을 생성하고, 이를 파라미터가 고정된 LLM 디코더에 넣어 원래 타겟을 복원하도록 학습한다. 오른쪽 패널에서는 S1S_1으로 쿼리와 양성 샘플의 조건부 확률 분포를 정렬하고, S2S_2로 양성 샘플 생성 확률은 높이면서 음성 샘플 생성 확률은 낮추는 DPO 스타일의 최적화를 수행한다. 인코더만 학습하고 디코더는 항상 고정(freeze)한다는 점이 설계상 중요한데, 디코더가 원래의 언어 모델링 능력을 유지해야 분포 기반 유사도 측정이 의미가 있기 때문이다.

4.2 Preliminary: 인스트럭션 기반 텍스트 임베딩

텍스트 qq와 인스트럭션 tt가 주어지면, 임베딩 eq,t=E(q,t)e_{q,t} = E(q, t)를 생성하여 인스트럭션 tt에 따라 관련 문서 dDd \in D를 매칭한다. 인스트럭션을 바꾸는 것만으로 검색, 유사도 측정 등 다양한 다운스트림 태스크에 적응할 수 있는 구조다.

4.3 정보 압축: 판별적 임베딩에서 생성적 임베딩으로

디코더 기반 LLM에서 마지막 토큰의 히든 스테이트를 임베딩으로 쓰면, 이 벡터는 첫 번째 출력 토큰의 시맨틱을 담고 있을 뿐 입력 전체를 인코딩하지 못한다. AutoRegEmbed는 컨텍스트 압축(Context Compression) 개념을 차용하여 이 문제를 해결한다.

구체적으로, 텍스트 q=(q1,,qn)q = (q_1, \ldots, q_n)과 인스트럭션 t=(t1,,tm)t = (t_1, \ldots, t_m) 뒤에 kk개의 압축 토큰 c=(c1,,ck)c = (c_1, \ldots, c_k)를 붙인다. 여기서 kn+mk \ll n + m이다 (실험에서는 k=5k = 5). 이 (q,t,c)(q, t, c)를 인코더 EE에 넣어 압축 임베딩 ec=(ec1,,eck)e_c = (e_{c_1}, \ldots, e_{c_k})를 생성한다. 그 다음 이 ece_c를 파라미터가 고정된 디코더 DD에 넣어 관련 문서 dd를 생성하도록 강제한다. 디코더는 원래 컨텍스트나 인스트럭션에 접근할 수 없고 오직 압축 변수만 볼 수 있기 때문에, 정보 병목(information bottleneck)이 형성되어 압축 변수가 글로벌 시맨틱을 포착하도록 유도된다.

최적화 목표는 다음과 같다:

LIC=tlogpθD(dtd<t,ec)\mathcal{L}_{IC} = -\sum_{t} \log p_{\theta_D}(d_t | d_{<t}, e_c)

여기서 θE\theta_E는 인코더 파라미터(학습 대상), θD\theta_D는 디코더 파라미터(고정)이다. 디코더가 압축 임베딩만 보고 타겟 문서를 복원해야 하므로, 인코더는 입력의 글로벌 시맨틱을 최대한 압축 토큰에 담도록 학습된다. 이렇게 생성된 임베딩은 조건부 확률 분포를 모델링하므로 **생성적 임베딩(Generative Embedding)**이라 부르며, 기존 풀링 방식의 판별적 임베딩과 구별된다.

4.4 조건부 분포 정렬: 데이터 포인트에서 분포 관점으로

글로벌 시맨틱 표현 문제를 해결했으니, 이제 임베딩이 정렬(Alignment)과 균일성(Uniformity)을 갖추도록 해야 한다. 전통적으로 이 두 속성은 InfoNCE 같은 대조 손실로 최적화한다:

LInfoNCE=logexp(sim(eq,ed+)/τ)exp(sim(eq,ed+)/τ)+iexp(sim(eq,edi)/τ)\mathcal{L}_{InfoNCE} = -\log \frac{\exp(\text{sim}(e_q, e_{d^+}) / \tau)}{\exp(\text{sim}(e_q, e_{d^+}) / \tau) + \sum_i \exp(\text{sim}(e_q, e_{d^-_i}) / \tau)}

여기서 τ\tau는 온도 파라미터이고 did^-_iii번째 음성 샘플이다. 이 수식은 코사인 거리를 최적화하므로 사전학습의 다음 토큰 확률 분포 모델링과는 근본적으로 다르다.

AutoRegEmbed는 코사인 거리 대신 조건부 확률 분포 간의 거리를 유사도로 쓴다. 디코더가 잘 학습된 자기회귀 언어 모델이라면, 임베딩 ece_c가 주어졌을 때의 조건부 확률 분포는:

pθD(dec)=t=1TpθD(dtd<t,ec)p_{\theta_D}(d | e_c) = \prod_{t=1}^{T} p_{\theta_D}(d_t | d_{<t}, e_c)

두 샘플 qqdd의 유사도를 코사인 거리가 아닌 조건부 확률 분포 간 거리 D(p(deq),p(ded))D(p(d|e_q), p(d|e_d))로 측정하는 것이 기본 아이디어다. 이를 InfoNCE 구조와 결합하여 최종 손실 함수를 경험적으로 도출한다:

L=logexp(σ(S1(q,d+))/τ)exp(σ(S1(q,d+))/τ)+iexp(σ(S2(d+,di;q))/τ)\mathcal{L} = -\log \frac{\exp(\sigma(S_1(q, d^+)) / \tau)}{\exp(\sigma(S_1(q, d^+)) / \tau) + \sum_i \exp(\sigma(S_2(d^+, d^-_i; q)) / \tau)}

여기서 S1S_1S2S_2는 각각:

  • S1(q,d+)=σ(βpθE(d+eq,Inext)pθE(d+ed+,Iself))S_1(q, d^+) = -\sigma\left(\beta \left| \frac{p_{\theta_E}(d^+ | e_{q, I_{next}})}{p_{\theta_E}(d^+ | e_{d^+, I_{self}})} \right|\right): 쿼리 임베딩 eq,Inexte_{q, I_{next}}와 양성 샘플 임베딩 ed+,Iselfe_{d^+, I_{self}}가 양성 샘플 d+d^+를 생성하는 로그 확률의 차이를 측정한다. 이 차이가 작을수록 두 분포가 유사한 것이므로, 절대값에 음수 부호를 붙여 유사도가 높을수록 S1S_1이 커지게 한다.

  • S2(d+,d;q)=σ(βpθE(d+eq,Inext)pref(d+eq,Inext)βpθE(deq,Inext)pref(deq,Inext))S_2(d^+, d^-; q) = -\sigma\left(\beta \frac{p_{\theta_E}(d^+ | e_{q, I_{next}})}{p_{ref}(d^+ | e_{q, I_{next}})} - \beta \frac{p_{\theta_E}(d^- | e_{q, I_{next}})}{p_{ref}(d^- | e_{q, I_{next}})}\right): DPO에서 영감을 받은 항으로, 쿼리 임베딩이 양성 샘플을 생성할 확률은 높이고 음성 샘플을 생성할 확률은 낮추는 차이를 계산한다. prefp_{ref}는 학습 전 초기 모델의 확률로, 양성·음성 샘플 길이 차이로 인한 편향을 보정한다.

σ()\sigma(\cdot)는 시그모이드 함수로, 조건부 확률 분포에서 측정한 유사도를 [0, 1] 범위로 정규화하여 지수 연산에서 오버플로를 방지한다. τ\tau는 InfoNCE의 온도 파라미터, β\beta는 DPO 스타일의 스케일링 파라미터다.

InextI_{next}IselfI_{self}는 각각 다른 인스트럭션을 가리킨다. 예를 들어 STS 태스크에서는 둘 다 "This sentence means in one word:"이고, 검색 태스크에서는 쿼리용과 문서용이 다르다. 추론 시에는 여전히 eq,Inexte_{q, I_{next}}ed,Iselfe_{d, I_{self}} 사이의 내적(inner product)으로 유사도를 계산하므로, 기존 임베딩 시스템과 호환된다.

4.5 학습 및 추론

학습은 두 단계로 나뉜다.

1단계 — 정보 압축: PWC 인스트럭션 파인튜닝 데이터셋에서 중복을 제거한 PWC-Unique (16,382 샘플)로 학습한다. 학습률 2e-5, 배치 크기 32, 2 에폭, 압축 토큰 5개. A100-80G 4대에서 약 20분 소요.

2단계 — 조건부 분포 정렬: MEDI의 NLI 파트 (50,000 샘플) 또는 BGE (274,951 샘플)를 사용한다. 학습률 5e-6, 배치 크기 32, 4 에폭, τ=0.05\tau = 0.05, β=0.1\beta = 0.1. 50,000 샘플 기준 약 1시간 소요.

인배치 음성 샘플(in-batch negative)이 필요 없다는 점이 메모리 효율 면에서 큰 장점이다. 전통적 대조 학습은 배치 크기 64에서 64×64=4,09664 \times 64 = 4,096개의 페어와이즈 비교가 필요하지만, AutoRegEmbed는 64개만 필요해서 비교 횟수가 64배 줄어든다.

검색 태스크의 경우, 쿼리-문서 사이의 비대칭성 때문에 조건부 확률 기반 최적화와 코사인 거리 기반 평가 메트릭 사이의 간극이 커진다. 이를 완화하기 위해 추가 1 에폭의 대조 학습 파인튜닝을 수행하고, NV-Embed를 이용한 하드 네거티브 마이닝(랭크 30~210 구간에서 7개 선택)을 적용한다.


5. 실험 결과

5.1 실험 설정

  • 데이터셋: STS12~STS17, STS22, STS-B, BIOSSES, SICK-R (10개 STS 데이터셋) + MS MARCO, NFCorpus, SCIDOCS (3개 검색 데이터셋)
  • 평가 지표: STS — 스피어만 상관계수 (×100), 검색 — nDCG@10
  • 베이스라인: (1) 학습 없는 모델: LLaMA2-L/M, Mistral-L/M, Echo, PromptEOL, MetaEOL, GenEOL (2) 비지도 대조 학습: LLM2Vec-LLaMA2/Mistral (3) 지도 대조 학습: NV-Embed, SFR-Embedding-2_R, gte-Qwen2-7B-instruct, LLM2Vec, 공정 비교용 자체 베이스라인
  • 구현: bfloat16, FlashAttention 2, DeepSpeed Zero-2, A100-80G 4대

5.2 주요 결과 — STS 태스크

MethodParamsBIOSSESSICK-RSTS12STS13STS14STS15STS16STS17STS22STS-BAvg.Vol.
학습 없음
LLaMA2-L7B63.2965.1045.2670.8356.6962.4863.2749.76-7.7660.4356.91 (10)0
LLaMA2-M7B65.9660.0144.7664.1348.6662.3363.1664.3527.5953.5058.67 (10)0
GenEOL-LLaMA2-ChatGPT7B-78.7170.7883.2877.7582.1080.45--79.8378.99 (7)0
비지도 대조 학습
LLM2Vec-LLaMA27B82.4171.7765.3979.2672.9882.7281.0286.7063.4778.3276.41 (10)~160K
LLM2Vec-Mistral7B83.2975.5567.6583.9076.9783.8081.9185.5865.9380.4278.50 (10)~160K
지도 대조 학습 (SOTA)
NV-Embed7.73B85.5982.8076.2286.3082.0987.2484.7787.4269.8586.1482.84 (10)1,054K
SFR-Embedding-2_R7B87.6077.0175.6782.4079.9385.8284.5088.9367.1083.6081.26 (10)~1,751K
gte-Qwen2-7B-instruct7.49B81.3779.1679.5388.9783.8788.4886.4988.7567.1686.8183.06 (10)~791M
LLM2Vec-Mistral (Supervised)7B85.2483.7078.8086.3784.0488.9987.2290.1967.6888.6584.01 (10)544K
공정 베이스라인
LLaMA2-inbatch-M (50K)7B78.0983.1777.1082.8280.5387.4084.4390.0264.5987.1881.53 (10)50K
LLaMA2-inbatch-M (275K)7B77.4382.2677.9584.9082.0687.2286.4388.2266.4286.1281.90 (10)274,951
AutoRegEmbed
AutoRegEmbed-LLaMA27B84.6581.4679.9886.3583.3389.2186.9187.6765.9086.9883.24 (10)50K (+16K)
AutoRegEmbed-Mistral7B86.8480.3278.9286.1883.2988.9886.7588.7764.5387.2483.18 (10)50K (+16K)
AutoRegEmbed-LLaMA2 (BGE)7B85.6283.8779.6087.3684.2989.4387.7289.4667.7887.9684.31 (10)275K (+16K)
AutoRegEmbed-Mistral (BGE)7B87.4883.9079.5687.6484.1189.5887.4689.8767.7788.4884.59 (10)275K (+16K)

Table 1 (원논문): STS 태스크 결과 (스피어만 상관계수 ×100). Vol.은 학습 트리플렛 수, 괄호 안은 정보 압축 단계의 인스트럭션 파인튜닝 데이터.

결과를 세 가지 축으로 분석할 수 있다.

학습 없는 모델 대비: AutoRegEmbed-LLaMA2의 10개 데이터셋 평균 83.24는 기저 모델 LLaMA2-L의 56.91보다 26% 이상 높다. 프롬프트 최적화 기반 최고 성능인 GenEOL (78.99, 7개 기준)과 비교해도, AutoRegEmbed가 7개 기준으로 84.89를 달성하여 약 6%p 앞선다. 프롬프트 최적화 방법은 13B 모델에서도 한계가 뚜렷하다는 점에서, 학습 기반 접근의 필요성을 확인할 수 있다.

비지도 대조 학습 대비: LLM2Vec-Mistral은 약 160,000개 샘플로 78.50을 달성했지만, AutoRegEmbed-LLaMA2는 50,000개(+16,382개) 샘플로 83.24를 달성했다. 4.74%p 차이를 데이터의 1/3로 만들어냈다는 건, 학습 효율성에서 확실한 우위가 있다는 뜻이다.

지도 대조 학습 대비: 50K 샘플 기준으로도 이미 NV-Embed (82.84, 1,054K 샘플)를 넘어서고, BGE 데이터 275K를 쓰면 84.59로 LLM2Vec-Mistral의 지도 학습 버전(84.01, 544K 샘플)도 앞선다. 동일 데이터로 공정 비교한 자체 베이스라인(LLaMA2-inbatch-M, 81.53)과는 1.71%p 차이를 보인다. SOTA 모델들이 수백만~수억 개의 멀티태스크 데이터를 쓴다는 점을 감안하면, 6만여 개의 데이터로 이 성능을 달성한 것은 주목할 만하다.

5.3 검색 태스크 결과

MethodMS MARCONFCorpusSCIDOCS
LLM2Vec (Unsupervised)18.8126.8110.00
LLM2Vec (Supervised)41.4540.3321.05
SFR-Embedding-2_R42.1841.3424.69
gte-Qwen2-7B-instruct45.9840.6023.48
LLaMA2-inbatch-M41.6734.1916.15
AutoRegEmbed42.4938.1619.79

Table 2 (원논문): MS MARCO, NFCorpus, SCIDOCS 검색 평가 결과 (nDCG@10).

MS MARCO에서 AutoRegEmbed는 42.49로 대부분의 SOTA 모델을 앞서고, gte-Qwen2-7B-instruct(45.98)에만 뒤진다. 동일 데이터로 학습한 LLaMA2-inbatch-M(41.67)보다는 세 데이터셋 모두에서 높은 성능을 보인다.

NFCorpus와 SCIDOCS에서는 SOTA에 미치지 못하는데, 이는 AutoRegEmbed가 MS MARCO만으로 학습된 반면 기존 SOTA 모델들은 다양한 도메인의 대규모 데이터셋으로 학습되었기 때문이다. 검색 태스크의 비대칭적 특성상 조건부 확률 기반 최적화와 코사인 거리 기반 평가 사이의 간극도 STS보다 크게 작용한다. 그럼에도 동일 조건에서의 일관된 우위는 방법의 유효성을 증명한다.

5.4 Ablation Study

MethodAvg. (10 datasets)
AutoRegEmbed-LLaMA283.24
w/o Conditional Distribution Alignment73.90
LLaMA2-L (Without Training)56.91
Equation 2 변형
Log_sigmoid82.93
KL divergence79.82
JS divergence79.02

Table 3 (원논문): AutoRegEmbed의 Ablation 실험. 태스크 제거 및 Equation 2 변형 비교.

태스크별 기여도가 명확하다. 조건부 분포 정렬을 제거하면 73.90으로 9.34%p 하락하고, 정보 압축까지 제거하면(즉, 학습 없는 LLaMA2-L) 56.91로 추가 16.99%p 하락한다. 정보 압축이 더 큰 기여를 하는데, 이는 글로벌 시맨틱 확보가 임베딩 품질의 기본 전제 조건이기 때문으로 보인다.

손실 함수 변형 실험에서는 원래의 Equation 2가 가장 좋은 성능을 보인다. Log_sigmoid 변형은 82.93으로 근소하게 뒤지고, KL/JS 다이버전스는 각각 79.82, 79.02로 크게 떨어진다. 저자들은 KL과 JS가 전체 어휘에 걸쳐 분포 차이를 측정하기 때문에, LLaMA2처럼 30,000개 이상의 어휘를 가진 모델에서 그래디언트 신호가 불안정해진다고 분석한다. 반면 log-odds ratio 방식은 특정 레이블 토큰의 생성 확률을 직접 감독하므로, 더 강하고 집중된 학습 신호를 제공한다는 것이다.

5.5 학습 효율성

Figure 1: AutoRegEmbed vs 기존 방법의 파레토 프론트
Figure 1: AutoRegEmbed vs 기존 방법의 파레토 프론트

Figure 1: AutoRegEmbed와 다른 방법들의 파레토 프론트 비교. 가로축은 학습 샘플 수, 세로축은 10개 STS 데이터셋 평균 성능. 왼쪽 위 영역이 학습 효율이 가장 높은 구간이다. (원논문)

Figure 1은 이 논문의 가장 인상적인 시각화다. AutoRegEmbed의 파레토 프론트(빨간 점선)가 대조 학습 파레토 프론트(회색 점선)를 완전히 지배한다. 특히 소량 데이터 영역에서의 격차가 크다. AutoRegEmbed는 5만 개 미만의 샘플로 이미 83점대에 진입하는 반면, 단일 태스크 대조 학습은 5만 개에서도 81점대에 머문다. 오른쪽의 SOTA 모델들(NV-Embed, SFR-Embedding-2_R, gte-Qwen2-7B-instruct 등)이 수백만~수억 개의 데이터(상당수 비공개)를 써서 도달한 성능에, AutoRegEmbed는 데이터의 극히 일부로 근접하거나 초과하고 있다.

Figure 3: 학습 효율성 비교
Figure 3: 학습 효율성 비교

Figure 3: AutoRegEmbed와 전통적 대조 학습의 10개 STS 데이터셋 학습 효율 비교. 동일 샘플 수에서의 성능을 비교한다. (원논문)

Figure 3은 더 세밀한 분석을 보여준다. 샘플 수를 점차 늘려가면서 성능 변화를 추적하면, AutoRegEmbed(파란 실선)가 모든 구간에서 네 가지 대조 학습 베이스라인을 앞선다. 15,000개 샘플만으로도 다른 대조 학습 모델의 최대 성능(약 80~81)을 넘어선다. "Ours consistently better"라는 그래프 상단의 점선이 이를 시각적으로 잘 보여준다. 인배치 네거티브를 쓰는 Lasttoken-InBatch와 Mean-InBatch는 초반에 빠르게 올라가지만, 이내 AutoRegEmbed에 추월당한다.


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

6.1 구현 세부사항 (Appendix A)

AutoRegEmbed의 정보 압축 단계는 학습률 2e-5, 배치 크기 32, 2 에폭으로 20분이면 끝난다. 조건부 분포 정렬은 학습률 5e-6, 배치 크기 32, 4 에폭에 50K 샘플 기준 약 1시간이다. 최대 토큰 길이는 컨텍스트, 인스트럭션, 타겟 모두 512다.

공정 비교를 위한 대조 학습 베이스라인은 FlagEmbedding 리포지토리 기반으로 구현되었다. 인배치 네거티브 전략을 쓰는 모델은 배치 크기 512가 최적이었고(128, 256, 512, 1024 중), 학습률은 1e-4가 최적이었다(1e-5, 5e-5, 1e-4, 2e-4 중). 그래디언트가 서로 다른 디바이스 간에 전파되도록 했다.

6.2 Equation 2 변형 상세 분석 (Appendix B)

Log_sigmoid 변형: InfoNCE의 로그 구조에 맞춰 시그모이드를 Log-Sigmoid 함수로 대체한다. 성능은 82.93으로 원래 Equation 2(83.24)보다 0.31%p 낮다.

KL 다이버전스 변형: 로그 확률 차이를 조건부 확률 분포 간의 KL 다이버전스로 대체한다. 성능은 79.82로 3.42%p 하락한다.

JS 다이버전스 변형: KL 대신 JS 다이버전스를 쓴다. 성능은 79.02로 4.22%p 하락한다.

KL/JS가 성능이 낮은 이유에 대해 저자들은 다음과 같이 분석한다. Log-odds ratio는 특정 레이블 토큰의 생성 확률을 직접 감독하여 강한 학습 신호를 제공하지만, KL/JS는 전체 어휘 분포에 걸쳐 작동하므로 대규모 어휘(~30K 토큰)에서 그래디언트가 불안정해진다. 이 분석은 직관적으로 타당한데, 분포의 대부분이 거의 0에 가까운 확률을 가진 토큰들로 구성되어 있어 불필요한 노이즈가 많아지기 때문이다.

6.3 σ 함수 선택 분석 (Appendix C)

MethodBIOSSESSICK-RSTS12STS13STS14STS15STS16STS17STS22STS-BAvg.
AutoRegEmbed-LLaMA2 (Tanh)83.9780.7580.5886.9283.1988.9886.9685.8065.9185.9884.77 (7) / 82.90 (10)
AutoRegEmbed-LLaMA2 (Sigmoid)85.5079.0779.5786.9083.2888.4586.5788.6166.1686.5984.35 (7) / 83.07 (10)

Table 4 (원논문): σ 함수(Tanh vs Sigmoid) 비교. 동일 설정에서의 성능.

Tanh가 7개 데이터셋 기준으로 84.77로 근소하게 앞서지만, 10개 기준으로는 Sigmoid가 83.07로 더 높다. 성능 차이가 0.17%p 이내로 미미하여, 어느 쪽을 선택해도 무방하다. 저자들은 Sigmoid를 기본값으로 채택했다.

6.4 InextI_{next}IselfI_{self} 인스트럭션 설명 (Appendix D)

태스크InextI_{next}IselfI_{self}
Retrieval"Use one word to represent the query in a retrieval task. The word is:""Use one word to represent the passage in a retrieval task. The word is:"
STS"This sentence means in one word:""This sentence means in one word:"

Table 5 (원논문): 검색 및 STS 태스크에서 사용하는 인스트럭션.

STS는 대칭적 태스크이므로 두 인스트럭션이 동일하고, 검색은 비대칭적이므로 쿼리용(InextI_{next})과 문서용(IselfI_{self})을 구분한다. PromptEOL의 선행 연구를 기반으로 한 이 인스트럭션 설계는 LLM이 입력의 의미를 한 단어로 압축하도록 유도하는 간단하면서도 효과적인 전략이다.

6.5 정렬 전략 분석 (Appendix E)

저자들은 기본 정렬 전략 외에 네 가지 대안 전략을 실험했다.

  • Strategy 1: d+d^+ 대신 qq를 생성 타겟으로 사용 (생성 대상 변경)
  • Strategy 2: qqd+d^+ 모두를 앵커로 사용하여 가중 유사도 계산
  • Strategy 3: 양성 샘플이 자기 자신을 생성하는 확률을 높이면서 음성 샘플 생성 확률을 낮추는 추가 항 도입
  • Strategy 4: 음성 샘플이 자기 자신을 생성하는 확률을 높이면서 양성 샘플 생성 확률을 낮추는 항 도입
MethodBIOSSESSICK-RSTS12STS13STS14STS15STS16STS17STS22STS-BAvg.
Strategy 186.3181.9781.0586.2782.8088.5186.3186.4364.2386.3584.75 (7) / 83.02 (10)
Strategy 283.3481.8779.8284.1181.9789.1186.3988.7966.1186.9084.32 (7) / 82.84 (10)
Strategy 386.0383.2078.3581.5481.2787.8087.0489.1664.5586.2883.64 (7) / 82.52 (10)
Strategy 484.1080.7478.0983.7981.7788.0186.1489.6464.1487.1683.67 (7) / 82.36 (10)
AutoRegEmbed85.5079.0779.5786.9083.2888.4586.5788.6166.1686.5984.35 (7) / 83.07 (10)

Table 6 (원논문): 서로 다른 정렬 전략의 STS 벤치마크 성능 비교.

일관된 패턴이 보인다: 원래 손실 함수에 가해지는 변형이 단순할수록 성능이 좋다. Strategy 1은 d+d^+ 대신 qq를 생성하도록 바꾼 것인데, STS의 대칭성 덕분에 큰 차이가 없다(83.02 vs 83.07). Strategy 2는 가중 평균이라는 추가 계산에도 불구하고 성능이 떨어진다. Strategy 3, 4는 양성/음성 샘플에 대한 더 복잡한 정렬을 설계했지만 오히려 성능이 하락한다. 이 결과는 복잡한 정렬 설계보다 단순하고 직접적인 감독 신호가 더 효과적이라는 점을 시사한다.

6.6 온도 계수 τ\tauβ\beta 분석 (Appendix F)

τ\tauβ\betaBIOSSESSICK-RSTS12STS13STS14STS15STS16STS17STS22STS-BAvg.
0.10.185.5079.0779.5786.9083.2888.4586.5788.6166.1686.5984.35 (7) / 83.07 (10)
0.020.184.9079.8578.5885.5484.6488.7186.8887.5765.6186.0284.32 (7) / 82.83 (10)
0.050.184.6581.4679.9886.3583.3389.2186.9187.6765.9086.9884.89 (7) / 83.24 (10)
0.20.183.8581.7279.7786.7383.1988.4186.5387.6966.2286.3784.67 (7) / 83.05 (10)
1.00.181.2380.5777.6383.9081.9287.0885.7588.1863.9185.9583.26 (7) / 81.61 (10)
0.10.284.0681.4980.3287.1583.4988.6786.8587.2266.4586.5684.93 (7) / 83.23 (10)
0.10.384.5079.5679.4586.6283.5587.7886.0189.7765.6386.1384.16 (7) / 82.90 (10)
0.10.484.2781.0478.9285.7682.4488.5486.0587.4865.9986.0884.12 (7) / 82.65 (10)
0.050.284.3179.6579.5984.1681.9589.5387.5489.3766.7887.4884.27 (7) / 83.04 (10)

Table 7 (원논문): 온도 τ\tau와 정렬 가중치 β\beta의 Ablation. Avg.는 7개 및 10개 데이터셋 기준.

τ\tau에 대해서는 0.05가 최적이고, β\beta에 대해서는 0.10.2 범위가 좋다. τ=1.0\tau = 1.0처럼 극단적인 값에서는 성능이 1.6%p 이상 떨어지지만, 나머지 설정에서는 비교적 안정적이다(82.6583.24 범위). 이는 방법의 로버스트함을 보여주는 긍정적 신호다.

τ=0.05\tau = 0.05β=0.2\beta = 0.2를 동시에 적용하면(83.04) 개별적으로 적용했을 때(각각 83.24, 83.23)보다 오히려 낮아지는데, 이는 두 온도 파라미터의 효과가 직교적이지 않다는 것을 시사한다. 저자들은 최종적으로 τ=0.05\tau = 0.05, β=0.1\beta = 0.1을 채택했다.


7. 강점과 한계

강점

  • 데이터 효율성: 66,382개의 학습 샘플(16K 정보 압축 + 50K 분포 정렬)로 수백만 개를 사용하는 SOTA 모델에 필적한다. Table 1에서 보듯, AutoRegEmbed-LLaMA2(83.24)가 1,054K 샘플의 NV-Embed(82.84)를 앞서는 것이 이를 입증한다.
  • 이론적 일관성: LLM의 자기회귀적 사전학습과 임베딩 학습 사이의 목표 불일치를 정면으로 다루면서, 코사인 유사도를 조건부 확률 분포 유사도로 대체한 것은 방법론적으로 설득력 있다.
  • 메모리 효율: 인배치 네거티브 없이 작동하여 페어와이즈 비교를 64배 줄인다. A100 4대에서 약 1시간 20분(압축 20분 + 정렬 1시간)이면 학습이 끝난다.
  • 기존 시스템 호환성: 추론 시 내적 기반 유사도를 사용하므로, 기존 벡터 검색 인프라와 바로 호환된다.

한계 및 아쉬운 점

  • 검색 태스크에서의 제한: STS에서는 강력하지만 검색 태스크에서는 추가 대조 학습 에폭이 필요하고, 그래도 단일 도메인(MS MARCO) 학습으로는 범용 검색 성능에 한계가 있다. Table 2에서 NFCorpus, SCIDOCS의 성능 격차가 이를 보여준다.
  • MTEB 전체 벤치마크 부재: 10개 STS 데이터셋과 3개 검색 데이터셋만 평가했다. MTEB에는 분류, 클러스터링, 재순위 등 다양한 태스크가 포함되어 있는데, 이런 태스크에서의 성능은 알 수 없다.
  • 베이스 모델의 제한: LLaMA2-7B와 Mistral-7B에서만 실험했다. 최신 LLM(LLaMA3, Qwen2 등)이나 더 큰 규모(13B, 70B)에서의 스케일링 특성은 검증되지 않았다.
  • 정보 압축 단계의 데이터 의존성: PWC-Unique 16,382개가 정보 압축에 쓰이는데, 이 데이터셋의 인스트럭션 다양성과 도메인 분포가 다른 도메인으로의 일반화에 얼마나 영향을 미치는지 분석이 없다.
  • 안전성 문제: 저자들도 밝혔듯이, AutoRegEmbed는 학습 데이터에 편향된 내용이 있으면 그대로 임베딩에 반영할 수 있고, 이를 필터링하는 메커니즘이 없다.

8. 마치며

AutoRegEmbed는 LLM을 텍스트 임베딩 모델로 변환할 때의 근본적 비효율, 즉 자기회귀적 사전학습과 판별적 대조 학습 사이의 괴리를 정면으로 해결하려는 시도다. 정보 압축으로 글로벌 시맨틱을 확보하고, 조건부 분포 정렬로 자기회귀 프레임워크 안에서 정렬과 균일성을 달성한다는 접근은 방법론적으로 깔끔하다.

개인적으로는 Figure 1의 파레토 프론트가 이 논문의 기여를 가장 잘 보여준다고 생각한다. 데이터 효율성 관점에서 기존 대조 학습을 명확하게 지배하는 결과는, LLM의 사전학습 능력을 "낭비하지 않는" 학습 방법의 가치를 실증한다. 다만 검색 태스크에서의 추가 대조 학습 의존성과 제한된 벤치마크 범위는 후속 연구에서 다뤄져야 할 부분이다. MTEB 전체 태스크에서의 검증, 최신 LLM에서의 스케일링 실험, 그리고 멀티태스크 학습과의 결합이 자연스러운 다음 단계가 될 것이다.


References

  • BehnamGhader et al. "LLM2Vec: Large Language Models are Secretly Powerful Text Encoders." CoRR, 2024.
  • Lee et al. "NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models." CoRR, 2024.
  • Li et al. "Llama2Vec: Unsupervised Adaptation of Large Language Models for Dense Retrieval." ACL 2024.
  • Jiang et al. "Scaling Sentence Embeddings with Large Language Models." EMNLP Findings 2024.
  • Ge et al. "In-Context Autoencoder for Context Compression in a Large Language Model." ICLR 2024.
  • Rafailov et al. "Direct Preference Optimization: Your Language Model is Secretly a Reward Model." NeurIPS 2023.
  • Wang and Isola. "Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere." ICML 2020.
  • Gao et al. "SimCSE: Simple Contrastive Learning of Sentence Embeddings." EMNLP 2021.
  • Wang et al. "Improving Text Embeddings with Large Language Models." ACL 2024.