Training LLMs to be Better Text Embedders through Bidirectional Reconstruction

https://arxiv.org/abs/2509.03020

Training LLMs to be Better Text Embedders through Bidirectional Reconstruction

 

https://github.com/LUMIA-Group/Anchor-Embedding

GitHub - LUMIA-Group/Anchor-Embedding: The official implementation of the paper "Training LLMs to be Better Text Embedders throu

깃 허브가 있다!!! 

 

방법은 단순합니다.

기존 LLM에게 EOS Token은 단순한 구분자다!

그러니까 Eos token으로 Pooling 방식을 활용할거면 이 EOS가 뭘 의미하는지 가르쳐줘야 한다!! 

위에 나온 방식처럼 EBQ2D는 Query를 Embedding으로 만든 뒤 이를 통해서 문서를 생성하도록 한다. 

EBD2Q는 Query와 Doc순서만 바꾼 것이다.

 

이 두가지 방식을 섞어서 사용하면 높은 성과를 나올 수 있도록 model을 튜닝할 수 있다고 하네요 ㄷㄷ

 

마지막으론 다른 모델들이 사용하는 것과 마찬가지로 Contrastive Learning을 진행해줍니다

 

GritLM도 이기네요..?

public data로 재현해서 그런가...

일단 단순하게 잘 이뤄낸 성능인 것 같습니다.

 

문제 상황	디코더 LLM의 [EOS] 출력 임베딩은 사전학습에서 주로 “구분자”로만 쓰여 문맥 의미/텍스트 간 정렬을 학습하지 못함 → 임베딩 태스크 성능 제한. 이를 해결하기 위해 [EOS] 표현 자체에 의미 정렬을 주입할 전처리 학습 단계가 필요.
제안 방법(개요)	2-단계 프레임워크: (Stage I) 양방향 복원 — 질의 [EOS]로 문서 생성(EBQ2D), 문서 [EOS]로 질의 생성(EBD2Q)을 교사강요 CE로 학습하여 [EOS] 임베딩을 의미 앵커로 강화 → (Stage II) 대조학습으로 표현 공간을 추가 정렬. 구조/마스킹 변경 없이 적용.
학습 데이터(공개)	두 단계 모두 E5 공개 데이터 사용. Stage I에서는 (query, positive) 쌍만 사용. Stage II에서는 데이터셋별 인스트럭션을 적용 (표 6 요약: NLI/FEVER/HotpotQA/MIRACL/MSMARCO/NQ/QuoraDup/SQuAD 등).
학습 설정 — Stage I(Full FT)	스텝 2000, LR 4e-5, 총배치 512(grad. acc.), 워밍업 300, EOS 풀링, gradient checkpointing, FlashAttention-2, max length 512. 목적함수: EBQ2D/EBD2Q의 CE 손실 합.
학습 설정 — Stage II(PEFT)	LoRA(r=16), 스텝 1000, LR 2e-4, 총배치 512, Stage I과 동일한 EOS 풀링/GC 설정 및 E5 인스트럭션 기반 대조학습.
모델/범위	디코더 전용 LLM 다수(LLaMA-3.1/3.2, Qwen2.5, Mistral; 1B~8B 규모)에 동일 프레임워크 적용.
실험 구성	MTEB 전범위 및 15-태스크 subset으로 평가(초기 수렴 분석용). Stage 조합(무감독/Stage I/Stage II/둘다) Ablation 수행.
주요 결과(전반)	제안법은 다양한 모델/스케일에서 일관된 향상. MTEB 평균에서 베이스라인(대조만) 대비 상승(예: 표 15 평균값 증가). 또한 초기 125 step 구간부터 성능이 더 높고 거의 수렴된 상태로 시작.
Ablation(Stage I 구성)	Baseline vs OnlyD2Q vs OnlyQ2D vs Anchor(둘 다) — MTEB 평균 60.99 → 61.38/61.62 → 62.24로 결합(α=0.2) 이 최고. α=0.2/0.5/0.8 간 성능 변화 작음(0.2 최상).
학습 전략 비교(효과 검증)	“스텝만 더 많이”가 아님: Stage II 3000 step(63.74) < 우리 방법(Stage I 2000 + Stage II 1000=65.19). 즉, 향상은 양방향 복원 단계 자체의 효과.
수렴/효율	Stage I을 거친 모델은 Stage II에서 초기 손실이 낮고(>10→<1), 더 빠르고 매끄럽게 수렴. 훈련 시간(8×A100 예시): 1B ≈4h/≈2h, 3B ≈12h/≈5h(Stage I/II).
기여 요약	(1) [EOS] 역할 불일치를 문제화, (2) EBQ2D/EBD2Q로 [EOS]에 의미 정렬 주입 후 대조학습, (3) 공개 데이터만으로 MTEB SOTA(공개 세팅 한정) 및 빠른 수렴 입증.
한계 / 향후 과제	(i) Stage I은 검색/재랭킹에 특히 유리 → 분류/클러스터링용 목적함수 보완 여지, (ii) 2-단계는 수렴을 돕지만 총 학습 오버헤드 증가, (iii) 영어 중심 세팅 → 다국어 확장 필요.

 

좋아요. 업로드하신 논문(Anchor Embedding: “Training LLMs to be Better Text Embedders through Bidirectional Reconstruction”)을 기반으로, 연구 문제–방법–실험–결과–기여–한계를 체계적으로 정리했습니다. 필요한 부분마다 원문 근거를 달았습니다.

한 줄 요약

디코더 LLM을 임베딩 모델로 더 잘 쓰기 위해, [EOS] 임베딩을 “앵커”로 삼아 질의↔문서를 서로 복원하게 만드는 양방향 생성(Stage I) 후, 표준 대조학습(Stage II) 을 수행하는 2-단계 학습 프레임워크를 제안하고, 공개 데이터(E5)만으로도 MTEB에서 일관된 성능 향상을 보였습니다.

---

1) 문제 상황(Why)

• 디코더 LLM의 [EOS] 출력 임베딩은 사전학습에서 단순 구분자로 쓰여 의미 정렬(semantic alignment)을 학습하지 못해, 임베딩 태스크에서 잠재력이 제한됩니다.
  → 해결책: [EOS] 임베딩을 의미 앵커(anchor) 로 만들기 위해, 질의→문서/문서→질의 양방향 복원을 먼저 시켜 해당 임베딩에 의미와 상호 정렬 정보를 주입. 이후 대조학습으로 표현 공간을 다듬습니다.

---

2) 핵심 아이디어(What)

• Stage I: 양방향 복원(Bidirectional Reconstruction)
  • EBQ2D: 질의 Q의 [EOS] 임베딩 e_Q를 프리픽스로 붙여 관련 문서 D를 교사강요(teacher forcing)로 생성하고, CE로 학습.
  • EBD2Q: 문서 D의 [EOS] 임베딩 e_D로 질의 Q를 생성하고 CE로 학습.
  • 두 손실의 가중합: LStage I=αLQ2D+(1−α)LD2QL_\text{Stage I}= \alpha L_{Q2D} + (1-\alpha)L_{D2Q}. 저자들은 α=0.2로 고정(민감하지 않음).
  • 알고리즘 1의 요약: (1) e_Q,e_D 추출 → (2) 임베딩/토큰 임베딩 준비 → (3) e_Q로 D 디코딩 → (4) CE(D̂,D) → (5) e_D로 Q 디코딩 → (6) CE(Q̂,Q) → (7) 가중합 반환.
• Stage II: 대조학습(InfoNCE)
  • E5 공개 코퍼스로 인배치+하드 네거티브를 포함한 InfoNCE, 코사인 유사도/온도 τ=0.05 사용.
• 임베딩 추출(양 단계 공통): [EOS] 풀링으로 문장 임베딩을 얻음.

---

3) 방법론(How) — 단계별(실행 절차)

1. 데이터 준비
   • Stage I: (query, positive) 쌍만 사용(E5 공개 부분).
   • Stage II: E5 각 서브셋에 태스크별 프롬프트/지시문 적용(표 6 요약).
2. Stage I(Full FT)
   • 스텝: 2000 / LR 4e-5 / 총 배치 512(grad acc) / 워ーム업 300 스텝 / 길이 512 / FlashAttn-2, gradient checkpointing, EOS pooling.
3. Stage II(LoRA FT)
   • 스텝: 1000 / LoRA rank 16, LR 2e-4, 총 배치 512, 나머지 설정 Stage I과 동일.
4. 모델 범위
   • LLaMA-3.2-1B/3B, Qwen2.5-1.5B, Mistral-7B, LLaMA-3.1-8B 등 디코더 전용 LLM에 적용.

---

4) 수식(핵심 손실)

• EBQ2D: LQ2D=−∑t=1mlog⁡PΘ(dt∣eQ,d<t)L_{Q2D} = -\sum_{t=1}^{m}\log P_\Theta(d_t \mid e_Q, d_{<t})
• EBD2Q: LD2Q=−∑t=1nlog⁡PΘ(qt∣eD,q<t)L_{D2Q} = -\sum_{t=1}^{n}\log P_\Theta(q_t \mid e_D, q_{<t})
• Stage I 합성: LStage I=αLQ2D+(1−α)LD2QL_\text{Stage I}=\alpha L_{Q2D}+(1-\alpha)L_{D2Q} (α=0.2 권장)
• Stage II(InfoNCE): 온도-스케일 코사인 유사도, τ=0.05\tau=0.05 고정.

---

5) 실험 설정(데이터·하이퍼)

• 학습 코퍼스(E5 공개 부분): ELI5(0.1), HotpotQA, FEVER, MIRACL, MS MARCO Passage(0.5)/Document(0.2), NQ, NLI, SQuAD, TriviaQA, Quora Dup(0.1), MrTyDi, DuReader, T2Ranking(0.5) 등.
• 지시문 프롬프트(표 6): NLI/FEVER/HotpotQA/MSMARCO/NQ/QuoraDup 등 태스크별 Retrieval 인스트럭션.
• 학습 시간(8×A100 80GB 기준): Stage I/II가 각각 1B≈4h/2h, 3B≈12h/5h, 8B≈45h/13h.

---

6) 결과(정량) — 핵심 포인트

• MTEB 전체(56 태스크) 평균: 모든 모델 크기에서 Anchor(우리 방법) 가 Stage II-only 대비 +1.1~+1.35p 수준의 일관된 향상(예: 1B 60.99→62.24, 3B 62.33→63.55, 8B 64.06→65.30). 카테고리별로 Retrieval / Re-ranking에서 향상이 뚜렷.
• 초기 수렴 속도: Stage I을 거치면 Stage II 초반 25~125 스텝 구간에서 기저선 대비 높은 점수로 시작하고 더 빠르게 수렴(1B/3B/8B 모두).
• 학습 전략 비교(동일 예산 대비): Stage II 3,000 스텝보다 Stage I(2,000)+Stage II(1,000) 가 더 높음(63.74 < 65.19). 단순 추가 스텝이 아니라 Stage I 자체의 효과임을 입증.
• Stage I 구성의 기여(표 3): OnlyD2Q/OnlyQ2D도 이득이지만 둘 다 결합(α=0.2) 이 평균 62.24로 최고. α는 0.2/0.5/0.8에서 민감도 낮음.
• BEIR 비교(효율성): Anchor는 LLaMA2Vec의 무감독 10k 스텝 대비 2k 스텝으로 더 높은 NDCG@10(56.40 → 58.07).

---

7) 기여(Contributions)

1. [EOS] 역할 불일치 문제를 지적하고, 2) EBQ2D/EBD2Q를 통해 [EOS] 임베딩을 의미 앵커로 만드는 새 학습 단계를 제안, 3) 공개 데이터만으로도 MTEB SOTA 수준을 달성(동일 조건 대비).

---

8) 한계 및 논의

• 산업용 비공개 대규모 모델들과의 직접 비교는 배제(공정성 문제·데이터/엔지니어링 파이프라인 비대칭). 학술적으로는 합리적이나, 실사용 상위권 모델과의 직접 수치 비교는 제공되지 않습니다.
• 적용 대상이 본문 실험에선 디코더 LLM에 한정(1B~8B). 인코더·하이브리드 구조 일반화는 향후 과제.

---

9) 실험 재현 체크리스트(Your HEGA/KURE 적용 팁)

• Stage I
  • 배치에서 (Q, D⁺)만 사용, 각각 [EOS] 풀링으로 e_Q,e_D 추출 → Q2D/D2Q를 교사강요로 생성 → CE 손실의 가중합(α≈0.2). 2000 스텝 / LR 4e-5 / 총배치 512 / 워럽 300 / len 512 / FlashAttn-2 / GC on.
• Stage II
  • InfoNCE(τ=0.05), 인배치+하드네거티브, LoRA r=16, 1000 스텝 / LR 2e-4 / 총배치 512. 프롬프트는 표 6 템플릿 참조.
• 권장 포인트:
  • Retrieval/재랭킹을 주목 타겟으로 삼을 때 Stage I → Stage II 순서를 유지(초기 수렴+최종 성능 모두 유리).
  • α 민감도 낮음(0.2 권장) → 튜닝 부담 적음.

---

10) 한 눈에 보는 표(연구 핵심 요약)

구분 내용

문제	[EOS] 임베딩이 사전학습에서 의미 정렬을 학습하지 못해 임베딩 태스크에 부적합.
아이디어	[EOS]를 앵커로 쓰는 양방향 복원(Stage I) + 대조학습(Stage II) 의 2-단계.
방법(수식)	LQ2DL_{Q2D}, LD2QL_{D2Q} CE, LStage I=αLQ2D+(1−α)LD2QL_\text{Stage I}=\alpha L_{Q2D}+(1-\alpha)L_{D2Q}, Stage II는 InfoNCE(τ=0.05).
데이터/모델	E5 공개 부분(여러 QA/IR 코퍼스), LLaMA/Qwen/Mistral(1B~8B), 지시문 프롬프트.
설정	Stage I 2k steps, LR 4e-5, bs 512, len 512 → Stage II 1k steps, LoRA r=16, LR 2e-4, bs 512.
주요 결과	MTEB 평균: 1B +1.25, 1.5B +1.35, 3B +1.22, 7B +1.12, 8B +1.24p. Retrieval/재랭킹 향상이 큼.
수렴/효율	Stage I 후 Stage II 초기부터 고득점·빠른 수렴. Stage II 3000보다 Stage I+II(2000+1000)가 더 좋음(63.74 < 65.19).
Ablation	OnlyQ2D/OnlyD2Q도 이득, 결합(α=0.2) 이 최고(62.24). α 민감도 낮음.
시간/비용	1B(4+2h), 3B(12+5h), 8B(45+13h) @8×A100-80GB.
기여	[EOS] 불일치 지적, 양방향 생성 도입, 공개 데이터만으로 SOTA급 향상.
한계	산업용 비공개 모델과 직접 비교는 배제(공정성 이슈).

---

11) 해설: 왜 잘 작동하나?

• Q→D 생성(OnlyQ2D) 는 짧은 질의로 긴 문서를 재구성해야 하므로 e_Q가 문서 수준 정보까지 담도록 강제 → Retrieval/재랭킹에서 특히 이득. 결합형(Anchor)이 평균 최고.
• Stage I 후 Stage II 는 이미 정렬된 표현 위에 밀도 조정만 하므로 빠른 수렴+높은 초기 성능.

---

 

맞아요, 핵심을 정확히 잡으셨어요.
요약하면:

• Stage I(양방향 복원):
  • eQ=EOS(Q)e_Q=\text{EOS}(Q)를 디코더의 연속 프롬프트로 넣어 문서 DD 를 교사강요로 생성(CE 손실).
  • eD=EOS(D)e_D=\text{EOS}(D)로 질의 QQ 를 생성(CE 손실).
  • 이 과정이 EOS 임베딩(e_Q, e_D) 에 “상호 의미 정렬”을 강하게 주입 → EOS가 문장 의미의 앵커가 됨.
• Stage II(대조학습):
  • InfoNCE(코사인 유사도+온도)로 같은 의미는 더 가깝게, 다른 의미는 멀게 밀도 조정.
  • Stage I로 이미 잘 정렬된 표현을 미세 조정하므로 초기 성능·수렴이 좋아지고 MTEB가 상승.

핵심 포인트

• 단순히 “임베딩을 주고 생성”이 아니라, “EOS 임베딩을 조건으로 Q↔D를 재구성” 하며 EOS 벡터 자체의 표현력과 의미 포착 능력을 키운 뒤, 대조학습으로 마무리합니다.
• Q→D는 짧은 질의로 긴 문서 정보를 압축해 담게 만들어 질의 임베딩을 풍부하게 하고, D→Q는 문서에서 핵심 질의 단서를 뽑게 해 상호보완적입니다. 둘을 가중합(보통 α≈0.2) 으로 묶을 때 가장 좋았다는 보고가 있어요.
• 향상 폭은 특히 Retrieval/재랭킹 계열에서 두드러졌고, InfoNCE만 했을 때보다 일관되게 우위였습니다.

실전에 적용 팁

• 인퍼런스 때도 EOS 풀링으로 문장 임베딩을 뽑는 전제를 유지하세요.
• Stage I은 교사강요 CE로 2k 스텝 내외, 이후 Stage II InfoNCE 1k 스텝 내외가 깔끔하게 먹힙니다.
• α, 온도(τ), 배치/인배치 네거티브가 민감 파라미터인데, 기본값(α≈0.2, τ≈0.05, 큰 배치) 이 무난합니다.

정리: Q/Doc의 EOS 임베딩을 조건으로 서로를 생성하게 만들어 EOS 표현력을 끌어올리고, 그다음 대조학습으로 공간을 다듬어 MTEB 성능을 올린 것이다—맞습니다.

 

좋아요. 업로드하신 논문(Anchor Embedding: “Training LLMs to be Better Text Embedders through Bidirectional Reconstruction”)의 관련 연구를 계열별로 정리하고, 핵심 차이를 통해 본 논문의 명확한 기여를 도출해 드릴게요.

큰 그림: “LLM을 임베딩으로” 하는 두 흐름

• (A) 대조학습 중심(contrastive-only) 또는 생성+대조 결합형: E5, GritLM, Echo Embeddings 등.
• (B) 디코더 LLM 재활용(Repurposing) 고도화: LLM2Vec(양방향 주의·MNTP·대조), LLaMA2Vec(무감독 적응 단계) 등.

이 논문은 [EOS] 임베딩을 ‘앵커’로 만들기 위한 생성적 사전 단계(Stage I) 를 도입해 Q↔D를 양방향 복원시키고, 그 다음 대조학습(Stage II) 로 마무리하는 2-단계 프레임워크를 제안합니다. 핵심은 “EOS 표현 자체에 의미 정렬을 주입”하는 것.

---

관련 연구 분류 & 차이점 요약

계열 대표 방법 핵심 아이디어/목표 학습 목적함수 구조 변경 데이터/세팅 이 논문(Anchor)와의 차이

대조학습(공개 데이터)	E5	대형 LLM이 합성한 다국어 텍스트 쌍으로 표준 대조학습	InfoNCE 기반 대조학습	아키텍처 유지	공개 부분만 비교	Anchor는 생성적 Stage I(Q→D, D→Q)을 넣어 EOS에 의미를 선주입 후 대조학습 수행.
생성+대조 결합	GritLM	표현(대조) + 생성(LM) 를 동시에 최적화(L=λRep·LRep+λGen·LGen)	대조 + LM	아키텍처 유지	지시문 형식 데이터	Anchor는 단계적(two-stage) 접근으로, 생성 목표를 EOS 임베딩에 조건화(Q↔D 재구성), 이후 대조로 공간 정렬.
AR-LM 임베딩 개선(Zero-shot 친화)	Echo Embeddings	입력 반복으로 양방향 컨텍스트 근사, 반복 토큰 임베딩 추출	학습 없이/경량	아키텍처 유지	제로샷 향상 강조	Anchor는 학습적으로 EOS를 재구성 신호로 강화(Q↔D CE 손실) → 대조학습 결합으로 지도 하 정렬.
LLM 재활용(주의·마스킹 변경)	LLM2Vec	양방향 주의 허용 + MNTP + 대조학습으로 디코더를 임베더화	MNTP + 대조	주의 마스킹 변경	공개 베이스라인 비교	Anchor는 주의 마스킹을 건드리지 않고, EOS-조건 생성(Q→D, D→Q)로 의미를 주입 후 대조.
LLM 재활용(무감독 적응)	LLaMA2Vec	무감독 적응 단계 10k 스텝(bs=256), 임베딩 성능 향상	비감독 적응	일부 설정 변경	BEIR로 보고	Anchor는 2k 스텝(Stage I) 만으로 더 높은 BEIR NDCG@10(56.40→58.07) 및 효율성 보고.
산업용 임베딩	NV-Embed 등	7B 디코더 + latent-attention pooling, causal mask 제거, 하드네거티브 지시문 튜닝	대조 중심	모듈 추가/마스킹 변경	대규모 사유 파이프라인	저자들은 공정성 문제로 직접 비교 생략(사유 데이터/엔지니어링 불일치). Anchor는 공개 데이터만으로 MTEB SOTA(공개 데이터 한정) 목표.

---

이 논문의 방법이 “무엇이 다른가?”

1. EOS 임베딩을 생성 조건으로 쓰는 양방향 복원(Stage I)
   질의 EOS로 문서 생성(Q→D), 문서 EOS로 질의 생성(D→Q) 을 교사강요 CE로 학습해 EOS 표현 자체에 문서-수준·의도-수준 의미를 주입합니다. 이어서 대조학습(Stage II) 로 밀도·경계를 다듬습니다.
2. 효율성과 수렴 이점
   같은 공개 데이터(E5)만으로도 Stage I→Stage II 구성은 대조만 학습 대비 초기부터 높은 점수·빠른 수렴을 보였고(25~125 step 분석), 전체 평균이 +1.1~+1.35p 향상했습니다.
3. 단일 구성요소의 기여 검증
   OnlyQ2D/OnlyD2Q도 이득이지만 결합(α=0.2) 이 최고(Avg 62.24). α=0.2/0.5/0.8 민감도도 낮게 보고.
4. 비교적 적은 단계로 BEIR 상 우위
   Stage I 2k 스텝으로 LLaMA2Vec 10k 대비 더 높은 NDCG@10(58.07). (효율+성능)

---

정량적 근거(발췌)

• MTEB 56태스크 평균: 1B/1.5B/3B/7B/8B 모두 Anchor가 베이스라인(대조만)보다 +1.12~+1.35p 향상. 카테고리별로 Retrieval·Re-ranking에서 특히 큼.
• Ablation: OnlyQ2D가 단독으론 가장 유리하지만, Q2D+D2Q 결합(α=0.2) 이 최종 평균 최고.
• Early-stage: Stage II 초반 25~125 step 구간에서 Anchor 초기점이 대조만 대비 높고 수렴 빠름(1B/3B/8B 공통).

---

이 논문의 명확한 기여(Contributions)

1. 문제 제기: 디코더 LLM에서 [EOS]는 원래 구분자 토큰일 뿐인데, 임베딩에서는 의미 앵커가 필요하다는 역할 불일치를 지적.
2. 방법 제안: [EOS] 임베딩을 조건으로 Q↔D 양방향 복원(EBQ2D/EBD2Q) 을 수행하는 생성적 Stage I + 대조적 Stage II의 2-단계 학습. 주의 마스킹/아키텍처 변경 없이 표현을 강화.
3. 효율·성능 동시 달성: 공개 데이터만으로 MTEB(56태스크) 에서 일관된 향상과 BEIR 효율 우위(2k vs 10k) 를 입증.

---

인용문구(초안) — “차이/기여”를 한 문장으로

• “Unlike LLM2Vec, which changes the attention direction and adds MNTP, we keep the decoder architecture intact and inject semantic alignment into the [EOS] embedding via bidirectional Q↔D reconstruction before contrastive fine-tuning, yielding consistently better MTEB/BEIR results with fewer steps on public data.”

필요하시면 위 표를 그대로 논문 Related Work 섹션용 텍스트(국/영문) 또는 LaTeX 표/도식으로 변환해 드릴게요.