왜 pooling이 문제가 되는가
BERT 계열 인코더는 토큰마다 벡터를 출력한다. 그런데 검색·클러스터링·분류에서는 문장 하나당 벡터 하나가 필요하다. 이 "토큰 시퀀스 → 단일 벡터" 축약을 담당하는 것이 pooling이다. 흔히 쓰는 두 방식이 [CLS] 토큰 벡터를 그대로 쓰는 방법과, 모든 토큰의 평균을 내는 mean pooling이다.
문제는 사전학습 목표에 따라 두 방식의 품질이 크게 갈린다는 점이다. 순수 MLM(BERT)으로만 학습한 모델의 CLS는 문장 의미를 압축하도록 훈련된 적이 없다. 그래서 CLS를 그대로 임베딩으로 쓰면 코사인 유사도가 거의 무너진다. 반대로 문장 목표(NLI, contrastive)로 파인튜닝된 모델은 pooling 방식을 정해두고 학습하므로, 추론 때 다른 방식을 쓰면 성능이 급락한다.
두 방식의 동작 차이
CLS pooling은 시퀀스 맨 앞 특수 토큰의 마지막 히든 스테이트만 취한다. mean pooling은 padding을 제외한 실제 토큰들의 평균이다. 여기서 실무에서 가장 자주 나오는 버그가 attention mask를 무시한 채 padding까지 평균에 포함시키는 것이다.
import torch
def mean_pooling(last_hidden, attention_mask):
mask = attention_mask.unsqueeze(-1).float() # (B, T, 1)
summed = (last_hidden * mask).sum(dim=1)
counts = mask.sum(dim=1).clamp(min=1e-9) # 0 나눗셈 방지
return summed / counts
def cls_pooling(last_hidden, attention_mask=None):
return last_hidden[:, 0] # [CLS] 위치
mask 처리를 빠뜨리면 짧은 문장일수록 padding 벡터가 평균을 왜곡한다. 배치 안에 길이 편차가 크면 같은 문장이라도 배치 구성에 따라 임베딩이 달라지는, 재현 불가능한 버그가 된다.
모델별 정답은 학습 방식이 정한다
선택 기준은 취향이 아니라 모델 카드다. 대표적인 조합은 다음과 같다.
| 모델 계열 | 권장 pooling | 비고 |
|---|---|---|
| BERT/RoBERTa (MLM만) | mean | CLS는 문장 의미 압축 미학습 |
| sentence-transformers (MiniLM 등) | 모델별 지정 | 대부분 mean, 일부 cls |
| BGE, E5 | cls | CLS로 contrastive 학습됨 |
| GTE, Nomic | mean | 학습 시 mean 고정 |
즉 "mean이 항상 낫다"는 명제는 틀렸다. BGE에 mean을 적용하면 오히려 나빠진다. 학습 때 쓴 pooling을 추론에서 동일하게 재현하는 것이 유일한 규칙이다.
설정으로 강제하기
파이프라인 여러 곳에서 임베딩을 만들면 pooling 방식이 조용히 어긋나기 쉽다. 설정에 명시하고 코드에서 분기시키는 편이 안전하다.
embedding:
model: BAAI/bge-base-en-v1.5
pooling: cls # 모델 카드 기준, 임의 변경 금지
normalize: true # 코사인 유사도용 L2 정규화
max_seq_len: 512
POOLERS = {"mean": mean_pooling, "cls": cls_pooling}
def encode(model, tok, texts, cfg):
enc = tok(texts, padding=True, truncation=True,
max_length=cfg["max_seq_len"], return_tensors="pt")
with torch.no_grad():
out = model(**enc).last_hidden_state
vec = POOLERS[cfg["pooling"]](out, enc["attention_mask"])
if cfg.get("normalize"):
vec = torch.nn.functional.normalize(vec, p=2, dim=1)
return vec
정규화와 유사도 지표의 정합성
pooling 다음으로 자주 어긋나는 것이 L2 정규화다. 코사인 유사도를 쓸 거면 정규화한 벡터를 저장해야 한다. 정규화된 벡터끼리는 내적(dot product)이 곧 코사인이므로, 벡터 DB에서는 cosine 대신 더 빠른 ip 메트릭을 쓸 수 있다. 단 인덱싱 때와 질의 때 정규화 여부가 달라지면 순위가 뒤틀린다.
주의점과 검증
운영 전 반드시 확인할 것들이다.
- 인덱싱과 질의가 동일한 pooling·정규화를 쓰는지. 배치 서버와 온라인 서버가 다른 코드 경로면 특히 위험하다.
- pooling을 바꾸면 기존 임베딩과 호환되지 않는다. 전체 재색인이 필요하며, 섞이면 검색 품질이 조용히 나빠진다.
- 정답 판단은 벤치가 한다. 자사 데이터 200~500쌍으로 Recall@k나 nDCG를 재고, mean과 cls를 실제로 비교한 뒤 고정하라.
- truncation으로 잘린 긴 문서는 mean이든 cls든 뒷부분 정보를 잃는다. 이때는 청킹이 pooling 선택보다 영향이 크다.
정리하면 mean과 cls는 우열 관계가 아니라 모델 학습 방식에 종속된 설정값이다. 모델 카드를 확인해 방식을 고정하고, attention mask와 정규화를 정확히 처리하고, 자사 데이터로 검증하는 세 가지가 실무의 핵심이다.