Attention, pooling 방법에 따른 llm-> encoder 성능 - Pooling And Attention What Are Effective Designs For LLM-Based Embedding Models

https://openreview.net/forum?id=CWAvMSNUqT

Pooling And Attention: What Are Effective Designs For LLM-Based...

 

ICLR 2025 Reject된 논문입니다.

그래도 여러 모델을 다양한 pooling 방법, attention 방버을 통해 test했고, 어떤 방식을 써야 할까에 조금은 도움이 된 논문입니다.

코드도 있으니 구현해보기 너무 좋을 것 같아요 

다 성능을 말하지만 전부 데이터도 다르고, attention, pooling 방식 모두가 다르다...

이렇게 되다 보니 상세한 비교가 불가능하다. 

깊은 레이어에서 비슷한 임베딩을 보여주고, 초반 레이어는 변화가 많은 것을 알 수 있다.

두 번째 figure에서 Layer에 따른 성능을 보여주는데 Layer가 깊어져도 큰 차이가 없는 것을 보아 낮은 Layer에서도 높은 성능을 낼 수 있음을 보여준다.

 

논문에서 제안한 pooling 방식인데 여기엔 큰 관심이 없어서... 

Casual Attention에서 EOS token pooling이 생각보다 낮은 성능을 가진게 좀 아쉽네요... 

 

 

문제 상황 (Motivation)	기존 LLM 기반 임베딩 모델들은 다양한 pooling/attention 전략을 사용하지만 서로 다른 데이터, 모델, 학습 세팅으로 인해 설계 효과를 정량적으로 비교하기 어려움. 또한 대부분 통계적 유의성 없이 보고됨.
연구 목표	동일한 base LLM과 동일한 학습 데이터를 사용하여 pooling과 attention 전략만 바꿔가며 임베딩 성능에 어떤 영향을 주는지 공정하게 실험하고 분석
제안 방법 (Method)	▶ 총 5가지 pooling/attention 조합 실험 ▶ 새로운 기법 제안: Multi-Layers Trainable Pooling ▶ 모든 실험에서 동일한 학습 데이터(Mistral-7B, 1.4M 쌍), 동일한 contrastive learning 사용
Multi-Layers Trainable Pooling	모든 layer의 hidden state를 사용하여 layer 가중치와 cross-attention을 통해 정보를 통합 → 기존 last-layer 기반보다 다양한 semantic 정보를 활용 가능
실험 디자인	▶ Mistral-7B, Qwen2-0.5B 사용 ▶ LoRA (r=16), contrastive learning ▶ 평가: MTEB 전체 benchmark (STS, Retrieval, Classification, Clustering 등) ▶ 통계 검정: Wilcoxon Signed Rank Test
실험 조합 (5개)	Model 1: EOS + Causal Model 2: Last-layer trainable + Causal Model 3: Multi-layer trainable + Causal Model 4: Last-layer trainable + Bi-attn Model 5: Multi-layer trainable + Bi-attn
주요 결과	✅ Model 5 (Multi-layer + Bi-attn): STS/Retrieval 최고 성능 ✅ Model 1 (EOS + Causal): Classification/Clustering에서 더 안정적 📉 Task 간 성능 편차 → pooling/attention은 task-dependent
추가 실험	분류/군집용 데이터 추가 시 Retrieval 성능 하락 / Clustering 성능 향상 발생 → 서로 상충하는 목적 함수를 가짐
기여 (Contributions)	🔹 동일 조건에서 pooling/attention 설계 효과 정량 분석 🔹 새로운 pooling 구조 제안 🔹 통계적 유의성까지 검정 🔹 NV-Embed 등 기존 모델과 비교 실험 제공 🔹 공개 코드: https://github.com/yixuantt/PoolingAndAttn
한계 (Limitations)	🔸 제안 방식은 STS/Retrieval에는 좋지만 다른 task엔 부적합 🔸 mean pooling, NV-Embed 완전 구현 비교 초기엔 빠짐 🔸 ablation 및 하이퍼파라미터 분석 부족 🔸 성능 향상이 작고 불균형하며 기술적 참신성이 약하다는 지적
Reviewer 평가 요약	🔸 실험 설계는 우수 🔸 성능과 기여는 제한적 🔸 기초 해석 부족, 성능 편차의 원인에 대한 분석이 약함 🔸 논문은 reject 되었지만 후속 연구 가능성은 높게 평가됨

 

논문 “Pooling and Attention: What Are Effective Designs for LLM-Based Embedding Models?”은 다양한 Pooling 및 Attention 설계가 LLM 기반 임베딩 성능에 미치는 영향을 통제 실험으로 분석하고, 새로운 기법인 Multi-Layers Trainable Pooling을 제안한 연구입니다. 이 논문은 ICLR 2025 또는 NeurIPS 2024와 같은 수준의 학회에 제출되었으나, Reject 되었습니다.

---

📌 Reject 사유 요약 

주요 사유	내용 요약
1. 기술적 기여 부족	“Multi-Layers Trainable Pooling”이 기존 NV-Embed와 유사하며, 구조적 참신성이 충분하지 않다는 평가.
2. 성능 향상 미약 & 불균형	제안 기법은 STS/Retrieval에는 향상되지만, Classification/Clustering에는 성능 저하가 발생.
3. 실험 설계 미흡	중요한 ablation 실험, 하이퍼파라미터 튜닝 분석이 빠짐. 예: cross-attention query 수, dim, layer weight 없이도 되는가 등
4. 분석 부족	task별 성능 편차의 원인에 대한 이론적 설명이나 해석이 부족하다는 지적.
5. 비교 실험 누락	Bi-directional + mean pooling, 또는 진짜 NV-Embed 스타일 baseline과의 직접적 비교가 부족

---

📋 Reviewer 별 주요 코멘트

🧾 Area Chair TPte (Reject 권고)

• 제안된 방법이 일관되게 좋은 결과를 내지 못하며, 성능 향상도 작고 task 간 편차가 큼
• 특히 기술적 novelty가 부족하고, classification 성능 하락 원인에 대한 이론적 설명 미비

---

🧾 Reviewer pdQZ (Score: 5 - Reject 경계)

• 장점: 실험 규모, 통계적 검정
• 단점: 성능 향상이 “놀랍지 않고 작다”, 제안 기법이 “크게 새롭지 않다”
• 질문: STS/Retrieval과 Clustering/Classification 사이 성능 차이의 근본 원인이 뭔지 설명 필요

---

🧾 Reviewer YfBn (Score: 3 - Reject)

• 제안 기법이 NV-Embed와의 차별점이 부족
• Mean Pooling + Bi-attn 조건을 실험에서 제외한 것은 문제
• layer weight의 역할과 attention 중첩 설명이 모호
• “왜 classification data를 추가했는데도 성능이 낮아졌는가?”에 대한 분석 필요

---

🧾 Reviewer S4qG (Score: 5 → 상향조정, 여전히 Reject)

• 장점: 좋은 ablation, pooling 설계 아이디어는 흥미로움
• 단점: 초기 학습 데이터가 retrieval 위주라 classification task에 편향
• 제안에 따라 classification/clustering 데이터를 추가한 실험을 수행한 점은 긍정적
• 그러나 성능은 여전히 하락하거나 미미 → 학습 전략, 데이터 믹스, loss function 개선 필요

---

🧾 Reviewer jm3o (Score: 5)

• 장점: 설계 공정성, 통계적 검정, 체계적인 설명 및 시각화
• 단점:
  • 임베딩 설계 자체의 “근본적 이해와 해석이 부족”
  • 모델 구성요소의 ablation이 부족 (Q matrix, attention dim 등)
  • 단순한 설계와의 비교 부재 (ex. last-layer pooling만 쓰는 경우)

---

🔧 저자 응답 및 보완 노력

응답 항목	대응 내용
성능 편차 원인 분석	STS/Retrieval과 Classification/Clustering은 목표 함수가 근본적으로 다르며, 하나의 설계로 모두 잡기 어려움. 다양한 task용 데이터를 추가 실험하여 trade-off 관찰
전체 MTEB 결과 보완	Pair classification, summarization 등 전체 결과 테이블 포함
NV-Embed 비교 실험 보완	Model 4를 NV-Embed 방식과 맞췄으며, 진짜 NV-Embed-style baseline도 별도로 추가
Task-dependent Q 실험	instruction 기반 query vector를 학습에 반영했지만, 오히려 성능 하락. 개선 여지 있음
Layer weight 역할	Layer weight는 positional 정보 + 중요도 학습을 위해 필요. attention으로는 position 간 구분 불가

---

📈 학문적 기여는 인정되었지만 Reject된 이유 요약

평가 항목	논문이 받은 평
문제 정의	적절하고 실용적인 문제 제기 (LLM 임베딩 설계 최적화)
실험 설계	공정한 조건 (동일 모델, 데이터), 대규모 실험 및 통계 검정 수행
기술적 참신성	NV-Embed와의 차별성이 약하며 “구조 확장”에 가깝다고 판단
성능 향상	일부 task에선 유의미하지만, 다른 task에서는 하락 혹은 미미
해석력	왜 task별로 다르게 작동하는지에 대한 이론적 분석 미흡
비교 실험 및 ablation	일부 핵심 조합/설계가 실험에서 빠짐 (특히 Bi-attn + Mean 등)

---

💡 향후 보완 방향 제안

1. Ablation 실험 강화
   → Q의 수, dim, layer weight 유무, attention head 수 등 구성요소의 영향 정량 분석 필요
2. Task-aware 구조 설계
   → STS, Retrieval, Classification의 다중 목적 학습 구조 설계 고려 (예: multitask loss)
3. Root-cause 분석 강화
   → Layer별 정보 차이에 따른 성능 차이 분석 (representation probing, attribution 등)
4. 간결한 baseline과 직접 비교
   → Multi-layer 구조 없이도 유사 성능이 나는 경우도 실험해야 설계 정당성 확보
5. Fine-grained task analysis
   → Retrieval 안에서도 entity vs. passage retrieval 성능 분석처럼 task 내 세분화 필요

---

 

 

---

📌 핵심 요약표

문제 상황	다양한 LLM 기반 임베딩 모델들이 서로 다른 학습 세팅·데이터·설계로 비교가 어렵고, Pooling 및 Attention 전략의 효과를 명확히 비교한 연구가 부족
연구 목표	동일한 LLM (Mistral-7B, Qwen2-0.5B), 동일 데이터, 동일 학습 방식으로 Pooling/Attention 전략만 다르게 실험하여 성능과 통계적 유의성을 분석
방법론	① Pooling 방식: EOS-last, Last-layer trainable, Multi-layer trainable ② Attention 방식: Causal vs. Bidirectional ③ 총 5가지 모델 조합으로 MTEB 벤치마크에서 비교
학습 데이터	총 1.4M 개의 쿼리-문서 쌍 (Quora, MSMARCO, SQuAD 등 18개), contrastive learning 사용
하이퍼파라미터	LoRA r=16, learning rate=1e-5, batch=2048, step=1000
결과 요약	- Multi-Layer + Bi-Attention은 STS/Retrieval에서 최고 성능 - Classification/Clustering은 Causal + Simple pooling이 더 효과적
기여	🔹 동등 조건 하 실험을 통한 설계 비교 🔹 새로운 Pooling 전략 제안 🔹통계적 검정을 통한 신뢰성 확보
한계	- Task에 따라 최적 구조가 다름 (one-size-fits-all 구조 없음) - Classification 등에서는 복잡한 구조가 오히려 성능 저하 가능성 있음

---

🧩 문제 상황 (Motivation)

• 기존 LLM 기반 임베딩 모델은 서로 다른 데이터, base LLM, pooling/attention 전략 사용 → 직접 비교가 불가능
• 성능 차이가 pooling 때문인지, attention 때문인지, 데이터 때문인지 명확하지 않음
• 대부분의 연구가 통계적 유의성(statistical significance)을 보고하지 않아 신뢰도 낮음

---

⚙️ 실험 구성 및 방법론

1. Pooling 전략

Pooling 종류	설명
EOS-last token	마지막 토큰의 hidden state 사용 (ex: E5, OpenAI text embedding)
Mean Pooling	전체 token 평균 (실험에선 제외: causal일 경우 성능 낮음)
Trainable Pooling (NV-Embed)	MLP 또는 Cross-Attn으로 hidden state 변환
Proposed: Multi-Layer Trainable Pooling	모든 layer의 hidden state를 cross-attention을 통해 조합하여 semantic space로 투영

2. Attention 전략

Causal	기존 LLM 학습 방식 (Unidirectional)
Bidirectional	Representation task에 적합. 모든 토큰이 서로를 볼 수 있음 (ex: GritLM, LLM2Vec 등)

---

🧪 실험 설정

• Base 모델: Mistral-7B-v0.1 및 Qwen2-0.5B
• 학습 방법: LoRA(16) + Contrastive learning (in-batch negatives)
• 총 5개 모델 조합:

모델	Pooling	Attention
Model 1	EOS-Last	Causal
Model 2	Last-layer Trainable	Causal
Model 3	Multi-layer Trainable	Causal
Model 4	Last-layer Trainable	Bi-directional
Model 5	Multi-layer Trainable	Bi-directional

• 평가: MTEB Benchmark (STS, Retrieval, Classification, Clustering 등 총 56개 task)
• 검정: Wilcoxon Signed Rank Test (p < 0.05 유의 기준)

---

📊 실험 결과 요약

✅ Pooling 전략별 성능

Task	EOS-Last vs Last-layer	EOS-Last vs Multi-layer	Last vs Multi-layer
STS	Multi-layer 우세 ★	Multi-layer 우세 ★	Bi-attn이면 Multi-layer가 우세 ★
Retrieval	유의하지 않음	Multi-layer 미세 우세	Bi-attn이면 Multi-layer가 우세 ★
Classification	단순 EOS-Last가 오히려 우수	복잡할수록 하락	Bi-attn에서 성능 저하
Clustering	단순 EOS-Last가 더 안정적	복잡할수록 하락	Bi-attn에서 Multi-layer는 일부 회복 ★

(★: 통계적 유의성 있음)

✅ 가장 우수한 조합

• Retrieval/STS: Model 5 (Multi-layer + Bidirectional)
• Classification/Clustering: Model 1 (EOS + Causal)

---

🧠 Multi-Layer Trainable Pooling 아키텍처 요약

• 입력: 모든 Layer의 Hidden States
• 가중치: Layer 별 가중치 W∈R^{l×d}
• Pooling 구조:
  1. Layer 가중치를 적용해 Layer Representation 생성
  2. Cross-attention block에 query Q는 trainable, key/value는 layer representation에서 생성
  3. Cross-attention 결과를 MLP로 투영하여 최종 embedding 생성

---

🧪 Robustness 실험: Qwen2-0.5B 사용

• 동일 조건으로 모델 1~5 학습
• 성능은 전반적으로 낮지만, 모델 간 경향은 Mistral과 일치
• STS/Retrieval에선 Multi-layer + Bidirectional 우세
• 작은 모델에선 복잡한 구조 도입해도 큰 이득 없음 → 단순 EOS + Causal도 효율적

---

📌 결론

• Task-specific 최적화가 필요: 하나의 구조로 모든 task 성능을 올릴 수는 없음
• Retrieval/RAG 등에서는 Multi-layer + Bi-Attn이 매우 효과적
• Classification/Clustering에는 단순한 구조가 더 낫기도 함
• 제안한 Multi-layer pooling은 정보 활용을 극대화하며, 향후 embedding 모델 설계에 중요한 전략으로 작용할 수 있음

---

 

논문 **“Pooling and Attention: What Are Effective Designs for LLM-Based Embedding Models?”**는 기존 LLM 기반 임베딩 모델 설계에서 pooling 및 attention 전략의 효과성을 공정하게 비교하고자 했으며, 이를 통해 기존 연구들 간의 모호한 비교 결과에 명확성을 제공하고자 했습니다.

---

🔍 1. Encoder 기반 임베딩 모델 연구 (Pre-LLM 시대)

이 계열은 BERT와 같은 encoder-only 모델을 중심으로 sentence embedding을 다루며, 비교적 단순한 구조에서 높은 효율성을 추구했습니다.

연구	모델	핵심 내용
BERT	BERT	Masked Language Modeling 기반의 사전학습. Encoder-only 구조.
Sentence-BERT (SBERT)	SBERT	Siamese 구조로 BERT를 사용하여 문장 임베딩 생성. Cosine similarity를 위한 contrastive 학습 도입.
INSTRUCTOR	T5-encoder 기반	Instruction을 입력에 포함하여 다양한 downstream task에 적용 가능한 universal embedding 모델 생성.
BGE-M3	XLM-R 기반	Dense + Sparse 임베딩을 결합하여 다양한 granularity를 학습하는 multi-function embedding model 제안.

➡️ 의의: 문장 임베딩 효율성을 획기적으로 높였지만, 대규모 LLM의 문맥 이해 능력을 충분히 활용하지 못함.

---

🤖 2. LLM 기반 임베딩 모델 연구 (본 논문의 직전 흐름)

대규모 LLM을 기반으로 임베딩을 수행하면서, 다양한 attention/pooling 전략을 도입한 선행 연구들이 있습니다.

연구	모델	Pooling 전략	Attention	주요 특징
RepLLaMA	LLaMA	Mean Pooling	Bi-directional	LLM을 retriever로 fine-tuning한 선구적 연구.
E5-Mistral	Mistral	EOS-Last Token	Causal	Instruction + Synthetic data를 활용한 contrastive 학습 기반 임베딩 모델
GritLM	Mistral	Mean Pooling	Bidirectional	Instruction tuning 기반. 일반 LLM보다 임베딩 성능 우수
LLM2Vec	LLaMA/Mistral	Mean Pooling	Bidirectional	LLM을 encoder로 활용하여, 임베딩 구조로 변형
NV-Embed	Mistral	Trainable Pooling	Bidirectional	Cross-attention 기반의 latent attention layer 도입. NVidia에서 제안한 고성능 임베딩 모델

➡️ 문제점: 서로 다른 학습 조건 (데이터, base LLM, pooling 구조 등)을 사용하여 설계의 효과를 직접 비교하기 어려움. 또한 대부분 통계적 유의성 없음.

---

🧠 3. 본 논문이 제안하는 핵심 기술과 연계된 연구

🧩 [1] Multi-layer Representation 활용 관련

• Clark (2019): BERT의 attention layer가 각기 다른 의미 정보를 담고 있다는 점을 분석
• Oh et al. (2022): 마지막 layer에만 의존하지 않고 layer-wise attention pooling을 통해 다양한 의미 정보를 추출
• Ju et al. (2024): LLM의 각 layer가 계층적인 의미를 다루는 것을 layer probing을 통해 확인

→ 본 논문의 "Multi-Layer Trainable Pooling"은 이러한 연구를 기반으로 실제 임베딩 모델에서 모든 layer의 정보를 trainable하게 결합하는 전략을 도입함.

🧩 [2] Cross-Attention 기반 Pooling

• Flamingo (Alayrac et al., 2022): cross-attention을 통해 다양한 길이의 vision frames를 고정된 latent embedding으로 변환
• NV-Embed (Lee et al., 2024): last layer의 hidden state를 cross-attention에 입력하여 고정된 dimension embedding으로 전환

→ 본 논문은 NV-Embed를 확장해 모든 layer의 hidden state + trainable layer weight + cross-attention을 결합하여 성능을 강화함.

---

📊 4. 평가 기준 및 벤치마크 관련 연구

연구	벤치마크	주요 내용
MTEB	Massive Text Embedding Benchmark	Retrieval, STS, Clustering, Classification 등 56개 task 포함
Springer et al. (2024)	임베딩 안정성 향상	같은 문장 반복을 통해 embedding 품질 개선 연구
Wang et al. (2023)	E5 시리즈	다양한 STS/IR task를 포함하여 모델 학습, 평가 및 MTEB 사용

→ 본 논문은 MTEB 전체 task에 대해 통계적 검정까지 수행한 점에서 기존 연구보다 신뢰성 있는 분석 제공.

---

🧩 기타 관련 기술

• LoRA (Hu et al., 2021): LLM의 parameter-efficient fine-tuning을 위한 low-rank adaptation 방법. 본 논문에서도 LLM embedding 학습에 사용.
• Contrastive Learning (Henderson et al., 2017): positive/negative 샘플 간 거리 최소/최대화를 통해 임베딩 학습. 본 논문은 in-batch negatives를 사용.

---

🔚 요약: 본 논문의 위치와 기여

측면	내용
기술적 기여	기존 연구의 단점을 보완하여 pooling 및 attention 전략의 효과를 통계적으로 정량 분석
설계적 기여	Multi-Layer Trainable Pooling이라는 새로운 pooling 구조 제안 (정보 손실 줄임)
벤치마크 기여	MTEB 전 task에 대해 동일 조건, 통계 검정 기반 성능 비교를 수행
한계 보완	서로 다른 데이터/모델 사용으로 인해 발생한 비교의 불공정성 제거

---

 

---

🔍 논문의 방법론 (Methodology)

본 논문의 핵심 목적은 다음과 같습니다:

❝ 동일한 LLM과 동일한 데이터로 pooling과 attention 전략만 바꿔서 임베딩 모델 성능에 어떤 영향을 주는지 정량적이고 통계적으로 분석하자 ❞

이 목표를 달성하기 위해 다음 3가지 설계 요소를 중심으로 연구를 설계했습니다:

---

🧩 1. Pooling 전략 (3가지 방식 비교)

LLM은 일반적으로 각 token에 대해 l개의 hidden layer 출력을 제공합니다.
이 때, 텍스트 전체를 대표할 고정 길이 임베딩 벡터를 만들기 위한 방법이 바로 Pooling입니다.

📌 1-1. EOS-Last Token Pooling (Baseline)

• 정의: 마지막 레이어의 마지막 토큰 (EOS)의 hidden state를 임베딩으로 사용
• 공식:
  h_{EOS} = H[-1, n-1, :]
  여기서 H는 (layers × tokens × hidden_dim) 형태의 hidden state
• 왜 사용하는가?
  causal LLM의 경우, 다음 토큰을 예측하도록 학습되기 때문에 마지막 토큰은 앞의 모든 정보를 요약한 representation 역할을 함.
• 예시:
• 입력: "How tall is Mount Everest? </s>" → </s> 토큰의 마지막 레이어 hidden state를 그대로 사용

---

📌 1-2. Last-Layer Trainable Pooling

• 정의: 마지막 레이어의 모든 토큰 hidden state를 입력으로 받아 trainable한 cross-attention 네트워크를 통해 임베딩을 생성
• 공식:
  h_{pool} = CrossAttention(Q, K, V)
  여기서 K, V는 마지막 layer의 hidden states로부터 생성, Q는 학습 가능한 query matrix
• 직관: 모델이 중요한 토큰에 집중하여 더 풍부한 임베딩을 학습할 수 있도록 유도
• 예시:
• 입력: "Albert Einstein was a theoretical physicist." → "Albert", "Einstein", "physicist"와 같은 중요한 단어를 attention으로 강조하여 임베딩 구성

---

📌 1-3. Multi-Layers Trainable Pooling (논문 제안 방식)

• 핵심 아이디어: 단일 레이어(hidden[-1]) 정보만 사용하지 말고 모든 레이어의 정보를 가중합 및 cross-attention으로 활용하자.
• 배경: LLM의 각 레이어는 다른 semantic 특성을 학습 (예: lower layer는 구조, 중간은 구문, 마지막은 task-specific)
• 구성 절차:
  1. 모든 레이어의 hidden state 수집
     • causal: EOS 토큰의 모든 레이어 hidden state → H[:, -1, :]
     • bidirectional: 각 레이어에 대해 token mean → 1/n * ∑_{i=1}^n H[:, i, :]
  2. Trainable Layer Weight 추가
     • 각 layer에 대한 중요도를 반영하기 위해 trainable matrix W∈R^{l × d}추가
     h_{layers} = H + W
  3. Cross-Attention Pooling
     • Key/Value: h_{layers}로부터 계산
     • Query: 학습 가능한 parameter Q∈R^{r×d'} 사용
     • Attention 후 MLP 통과 → 최종 임베딩 출력
• 장점:
  • 단일 레이어의 bias 제거
  • 다양한 semantic 정보를 통합할 수 있음
• 예시:
• 입력: "Photosynthesis is the process used by plants to convert light into energy." → 3~5 layer는 '과학적 개념'을, 20~24 layer는 문장 구조를 encode할 수 있음 → 이를 모두 조합하여 의미적으로 풍부한 embedding 생성

---

🧠 2. Attention 전략 (2가지 방식 비교)

📌 2-1. Causal Attention

• 정의: 현재 토큰은 과거 토큰만 바라볼 수 있음 (decoder LLM의 기본 방식)
• 특징:
  • 자연스러운 언어 생성을 유도
  • 하지만 임베딩 관점에서는 문맥 정보가 제한됨 (편향 발생)

📌 2-2. Bidirectional Attention

• 정의: 모든 토큰이 모든 토큰을 볼 수 있음 (BERT와 같은 encoder 구조)
• 장점:
  • 전체 문맥을 고려한 임베딩 가능
  • STS, IR task에서는 더 효과적
• 주의:
  • EOS-last token pooling과는 호환이 좋지 않음 (EOS가 "마지막"이라는 의미가 없음)

---

🧪 3. 학습 설정

Base LLM	Mistral-7B-v0.1
학습 방식	Contrastive Learning (positive/negative pair)
Parameter Efficient Fine-Tuning	LoRA (r=16)
데이터	총 1.4M 쌍 (MSMARCO, Quora 등 18개 open set)
Pooling 구조 설정	Trainable cross-attention: 32 heads, dim=4096
Evaluation	MTEB 전체 benchmark 사용 (STS, IR, Classification, Clustering 등 총 56 task)
통계 분석	Wilcoxon Signed Rank Test (p < 0.05 기준)

---

✅ 총 5개 모델 설계 비교

모델	Pooling	Attention
Model 1	EOS-Last	Causal
Model 2	Last-Layer Trainable	Causal
Model 3	Multi-Layers Trainable	Causal
Model 4	Last-Layer Trainable	Bidirectional
Model 5	Multi-Layers Trainable	Bidirectional

---

🎯 핵심 요약 

전략	효과적인 Task	비고
EOS-last + Causal	Classification, Clustering	단순하지만 안정적
Multi-layer + Bi-Attn	Retrieval, STS	성능 최상, 비용 높음
Last-layer Trainable	중간 수준	일부 task에서는 Overfitting 가능

---

💡 전문가 팁: 실전 적용 가이드

상황	추천 구조
RAG 시스템의 retriever	Model 5 (Multi-layer + Bi-attn)
단순 카테고리 분류/군집화	Model 1 (EOS + Causal)
작은 모델 기반 경량 임베딩	Model 1 또는 2 (Qwen2 등 소형 모델에 적합)
학습 시간/자원이 제한됨	Last-layer pooling + Causal로 시작

---

 

---

✅ 1. 결과 (Results)

논문에서는 동일한 LLM과 데이터셋을 기반으로 총 5가지 Pooling & Attention 조합을 실험하여 MTEB 벤치마크의 주요 task들(STS, Retrieval, Classification, Clustering)에 대해 성능을 비교했습니다.

📊 주요 결과 요약표 (Mistral-7B 기준)

모델	Pooling	Attention	STS	Retrieval	Classification	Clustering
Model 1	EOS-Last	Causal	✅ baseline	baseline	✅ 최고	✅ 최고
Model 2	Last-layer Trainable	Causal	▲ +0.0129*	▲ +0.0102	▼	▼
Model 3	Multi-layer Trainable	Causal	▲ +0.0118*	▲ +0.0135	▼	△
Model 4	Last-layer Trainable	Bi-directional	▲ +0.0095*	▲ +0.0213	▼▼	▼▼
Model 5	Multi-layer Trainable	Bi-directional	✅ 최고 ▲ +0.0166*	✅ 최고 ▲ +0.0226*	▼▼▼	▼▼

* = 통계적으로 유의미한 차이 (Wilcoxon p < 0.05)

🔍 세부 분석 결과

• Retrieval & STS:
  → Model 5 (Multi-layer + Bi-attention)이 가장 성능이 뛰어났으며, 기존 설계(Model 1) 대비 최대 +4.2% 향상.
• Classification & Clustering:
  → Model 1 (EOS + Causal)이 여전히 가장 강력하며, 복잡한 pooling/attention 전략은 오히려 성능 저하 유발 가능성 있음.
• Qwen2-0.5B 모델로 진행한 robustness 실험에서도 동일한 경향 유지됨.

---

🧠 2. 결론 (Conclusion)

논문의 핵심 결론은 다음과 같습니다:

❝ LLM 기반 임베딩 모델 설계에서 pooling과 attention 전략은 task에 따라 최적 조합이 달라진다. Retrieval 및 STS에선 복잡한 설계가 효과적이지만, classification/clustering에선 단순 설계가 오히려 더 나은 결과를 낳는다. ❞

📌 구체적 결론

1. Retrieval / STS
   → Multi-Layer Trainable Pooling + Bidirectional Attention 구조가 성능 및 통계적으로 가장 우수
   → LLM의 모든 레이어 정보를 활용해 semantic representation 강화
2. Classification / Clustering
   → 복잡한 구조는 오히려 정보 과다 및 과적합 위험
   → EOS-Last Token Pooling + Causal Attention이 더 안정적이며 효율적
3. 전반적인 시사점
   • “One-size-fits-all” 설계는 존재하지 않음
   • 실제 사용 시에는 task-specific 설계 전략이 필수

---

⚠️ 3. 한계 (Limitations)

논문은 강력한 실험 설계에도 불구하고 다음과 같은 한계를 가지고 있습니다:

항목	내용
제한된 LLM 베이스	Mistral-7B, Qwen2-0.5B만 사용 → GPT, Claude, Mixtral, LLaMA-3 등 대규모 모델에 대한 일반화 한계
Mean Pooling 제외	실험 스코프를 제한하기 위해 mean pooling은 제외되었지만, 일부 task에서는 여전히 유용할 수 있음
연산량 증가	Multi-layer pooling은 compute/resource 측면에서 훨씬 무거움
Downstream fine-tuning 미포함	task-specific fine-tuning과의 interaction은 분석하지 않음

---

🌟 4. 기여 (Contributions)

논문의 주요 학문적/실용적 기여는 다음과 같습니다:

카테고리	기여 내용
실험 설계	기존 연구들과 달리 동일한 LLM, 동일한 데이터, 동일한 학습 방식을 사용해 pooling/attention의 영향만 정밀 비교
새로운 기법 제안	모든 레이어 정보를 활용하는 Multi-Layers Trainable Pooling 제안 및 구현
통계적 검정 도입	Wilcoxon Signed Rank Test를 활용해 성능 차이에 대한 통계적 유의성 확보
공개 코드 제공	모든 실험 및 pooling module을 공개하여 reproducibility 확보
➡ https://github.com/yixuantt/PoolingAndAttn
실용적 인사이트 제공	실무 환경에서 task별로 어떤 설계가 적합한지 명확하게 가이드 제시 (RAG vs Clustering 등)

 

 

좋습니다. 논문에서 제안한 Multi-Layers Trainable Pooling은 이 연구의 핵심 기여 중 하나로, 기존의 단일 레이어 기반 임베딩(Pooling) 전략이 가진 정보 손실 문제를 해결하려는 의도로 설계되었습니다.

---

🧠 Multi-Layers Trainable Pooling: 자세한 설명

1. ❓ 왜 필요한가? (Motivation)

기존 대부분의 임베딩 전략은 LLM의 마지막 레이어(hidden[-1])만을 사용해서 임베딩을 구성합니다.

하지만,

• LLM의 각 레이어는 서로 다른 수준의 의미 정보를 담고 있음
  • 하위 레이어: 문법, 구문
  • 중간 레이어: 의미적 패턴
  • 상위 레이어: task-specific 정보
• 따라서 마지막 레이어만 사용하면 중간/하위 레이어의 풍부한 정보가 소실됨

☝️ 해결 아이디어
➜ "모든 레이어의 정보를 모아서 더 정교한 임베딩을 만들자."

---

2. 🧩 구조 요약 (설계 개요)

요약:
“모든 레이어의 출력(H)을 입력으로 받아서,
cross-attention 네트워크로 중요한 layer 정보를 추출하고,
이를 임베딩 벡터로 투영하는 학습 가능한 pooling 레이어”

---

3. ⚙️ 작동 절차 (Step-by-step)

🔸 Step 1: 모든 레이어의 hidden state 수집

• 입력 텍스트를 LLM에 넣으면 hidden states H는 다음과 같은 형태:
  H∈R^{L×N×D}
  • L: 레이어 수 (예: 32)
  • N: 토큰 수 (예: 128)
  • D: hidden dim (예: 4096)
• Causal Attention 모델이면 → 각 레이어에서 EOS token만 추출:
  h_{causal} = H[ : , −1, : ] ∈ R^{L×D}
• Bi-directional 모델이면 → 각 레이어의 mean vector 사용:
  h_{bi} = 1/N * ∑_{i=1}^N H[ : , i, : ] ∈ R^{L×D}

---

🔸 Step 2: Trainable Layer Weight Matrix 추가

• 모든 레이어가 동일하게 중요한 건 아님 → Layer마다 중요도를 학습
  h_{layers} = h + W
  • W ∈ R^{L×D}: 학습 가능한 layer weight
  • h: 위에서 만든 h_{causal} 또는 h_{bi}

---

🔸 Step 3: Cross-Attention을 통한 정보 집약

• 이제 hlayersh_{layers}를 Key, Value로 사용
  • K , V=h_{layers} ⋅ W_K, h_{layers} ⋅ W_V
• 학습 가능한 Query 벡터 Q ∈ R^{r×d′}사용
  (보통 r = 1, 즉 고정된 query 하나 → 1개의 임베딩 생성)
• Cross-Attention 연산 수행
• 결과를 MLP 통과시켜 최종 embedding 출력

---

🔸 Step 4: MLP 투영

• 위 attention 출력을 MLP로 정제:
• FinalEmbedding=MLP(CrossAttentionOutput)
• 결과적으로 shape: R^D 크기의 텍스트 임베딩 벡터가 완성

---

4. 🧪 예시: “What is quantum entanglement?”

1. 입력 텍스트를 Mistral-7B에 넣음 → 32개 레이어의 hidden state 생성
2. 각 레이어의 EOS token hidden vector만 추출 (causal 모델 기준)
3. 레이어별로 가중치를 다르게 부여하여 조합
4. cross-attention을 통해 가장 의미 있는 layer 정보 추출
5. 그 정보를 기반으로 고정 차원의 임베딩 생성

---

5. 📈 장점 vs 단점

장점	설명
🧠 정보 손실 방지	다양한 semantic depth의 정보를 반영 가능
🎯 task-specific 표현 가능	Layer 가중치를 학습함으로써 task에 따라 중요한 layer 강조
🔍 분석 가능성 증가	어떤 레이어가 중요한지 해석 가능 (Layer Attribution)

 

단점	설명
🧮 연산량 증가	모든 레이어를 저장하고 처리해야 함
💾 메모리 부담	특히 7B~13B 모델에서는 hidden state가 매우 큼
🐢 느린 inference	실시간 임베딩 추출에는 다소 무거울 수 있음

---

✅ 정리 요약

항목	내용
이름	Multi-Layers Trainable Pooling
목적	LLM의 모든 레이어 정보를 통합해 더 정밀한 임베딩 생성
핵심 구성	layer-weighted hidden state + cross-attention + MLP
최적 task	Semantic Textual Similarity, Information Retrieval
비교 우위	기존의 last-layer pooling보다 성능 ↑ (통계적으로 유의함)

---

🔧 참고 구현 스니펫 (PyTorch 스타일 예시)

class MultiLayerPooling(nn.Module):
    def __init__(self, num_layers=32, hidden_size=4096):
        super().__init__()
        self.layer_weights = nn.Parameter(torch.randn(num_layers, hidden_size))
        self.query = nn.Parameter(torch.randn(1, hidden_size))
        self.attn = nn.MultiheadAttention(embed_dim=hidden_size, num_heads=8, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(hidden_size, hidden_size), nn.ReLU())

    def forward(self, hidden_states):  # shape: [batch, layers, hidden]
        h = hidden_states + self.layer_weights  # broadcasting
        attn_out, _ = self.attn(self.query.unsqueeze(0), h, h)
        return self.mlp(attn_out.squeeze(0))

---