UltraBERT가 마스킹 비율을 분리해 분류와 검색을 동시에 최적화한다

무엇인가

양방향 인코더는 생성 모델이 잘하지 못하는 자리에 남아 있다. 추론 비용이 낮아 대량 텍스트를 임베딩하거나 점수화할 수 있기 때문에 검색증강생성(RAG) 파이프라인, 사전학습 데이터 품질 필터링, 강화학습용 보상 모델링의 백본으로 쓰인다. 최근에는 RoPE, FlashAttention-2, GLU, 어텐션 교대 배치, 조 단위 토큰 학습 같은 디코더 LLM의 혁신을 받아들이며 현대화됐다. 그런데 정작 인코더 고유의 마스크 언어모델링(MLM) 목표와 마스킹 비율은 BERT 시대 소규모 관행을 그대로 물려받았다. 그 결과 어떤 모델도 분류와 검색을 동시에 잘하지 못하고, 한쪽에 특화되는 상태가 이어졌다.

어떻게 동작하나

이 논문은 먼저 사전학습 레시피 자체를 조 단위 토큰 규모에서 다시 검토해 세 가지를 확인한다. 첫째, 150M과 1.3B 인코더를 8.8T 토큰까지 학습시키면 다운스트림 성능은 1~2T 토큰 부근에서 포화되는데, 학습 손실과 검증 손실은 계속 줄어든다. BEIR 점수는 275B 토큰의 첫 평가 지점 이후 사실상 개선이 없다. 같은 조건의 디코더(SmolLM2 1.7B)가 지시 튜닝 후에도 계속 향상되는 것과 대비되며, 저자들은 인코더에 추가 토큰으로는 올릴 수 없는 상한이 있다고 본다. 둘째, 포화 시점에서는 수렴 속도가 교란 변수로 작용하지 않는데, 이때 분류는 모든 모델 크기에서 낮은 마스킹 비율을 선호하고 검색은 50%가 모든 스케일에서 최고였다. 큰 모델에는 큰 비율이 좋다는 기존 연구(10B~100B 토큰 예산)와 반대 방향이다. 셋째, 프루닝 후 증류는 15% 마스킹에서 GLUE를 90.2에서 90.8로 올렸지만, 50% 마스킹에서는 BEIR를 47.2에서 46.5로 떨어뜨렸다.

무엇과 다른가

제안 방법의 핵심은 두 마스킹 체제를 한 모델 안에서 화해시키는 것이다. 저자들은 6%와 50% 마스킹으로 학습한 모델의 입력 임베딩은 거의 같지만 LM 헤드의 기하가 뚜렷이 달라진다는 점을 관찰한다. 6% 마스킹의 LM 헤드는 특이값이 더 집중돼 저차원 부분공간을 차지하고, 이는 표현 공간의 낮은 effective rank와 일치한다. 그래서 배치 안에 5%와 50% 마스킹 시퀀스를 4대 1로 섞어 넣고, 각 비율에 별도의 LM 헤드를 배정하는 LM 헤드 분리를 제안한다. 5%는 그래디언트 신호가 희소해 수렴이 느리므로 배치에 네 배 더 많이 넣는다. 증류는 로짓 기반 KL 발산 손실을 5% 시퀀스에만 적용하고, 50% 시퀀스는 표준 MLM으로 학습한다. 1.3B 교사 모델에서 출발해 380M, 다시 150M으로 프루닝한 뒤 1.3B 교사로부터 증류하며, 2T 토큰과 네이티브 32K 컨텍스트로 학습한다.

어떻게 쓰나

절제 실험 수치가 방법 선택의 근거를 뒷받침한다. 마스킹 비율을 50%에서 6%나 15%로 감쇠시키는 방식은 GLUE에서 단일 비율보다 나았던 적이 없고 6%에서는 오히려 88.8로 89.5보다 나빴다. 5~50%에서 균등 샘플링(평균 27.5%)은 단일 30%와 거의 같은 성능을 냈다( GLUE 88.0 대 88.3, BEIR 44.2 동일). 반면 5%와 50%를 섞고 LM 헤드를 분리하면 GLUE 89.4, BEIR 45.1로 각각 단일 6%(89.5)와 단일 50%(45.0)의 최고치에 도달했다. 헤드를 공유하면 GLUE 88.9, BEIR 44.7에 그쳤다. KoLeo 정규화로 effective rank를 인위적으로 높이는 실험도 시도했지만, BEIR은 개선되면서도 단일 50%보다 못했고 GLUE는 손상됐다. 참고로 Replaced Token Detection이 MLM보다 분류에 유리한 폭은 15% 마스킹 부근에서 최대로 벌어지고 양 끝으로 갈수록 줄어든다.

전제와 한계

최종 결과로 UltraBERT는 GLUE 91.0, BEIR 47.7을 기록하며 분류와 검색 양쪽에서 동시에 최고 성능을 낸 최초의 인코더라고 저자들은 주장한다. 150M 모델은 RoBERTa-large와 분류 성능이 대등하면서 파라미터가 2.3배 적고, Ettin-encoder-400M보다 검색이 우수하면서 크기가 2.6배 작다. 32K 컨텍스트를 기본 지원해 LongEmbed에서 두 크기 모두 최고 점수를 냈고, 코드·수학 검색(RAR-b)에서 2위 대비 base +7.3점·+8.3점, medium +8.4점·+7.7점으로 가장 큰 개선을 보였다. 예외도 있다. base 크기에서는 MultiCoNER v2에서 1위가 아니고, medium 크기에서는 MLDR에서 LFM2.5-Encoder-350M에 뒤진다. 비교 대상은 BERT, RoBERTa, MosaicBERT, NomicBERT, GTE-en-MLM, ModernBERT, Ettin Encoder, NeoBERT, LFM2.5-Encoder다.

실무적으로 이 논문이 주는 신호는 명확하다. 검색과 분류를 각각 다른 모델로 운영하던 팀이라면, 150M이나 380M 하나로 인덱싱·리랭킹·분류를 함께 처리할 수 있다는 선택지를 검토할 수 있다. 다만 그대로 가져다 쓰기 전에 확인할 것이 있다. 저자들의 주장은 인코더의 특화가 본질이 아니라 레시피의 결과라는 것이므로, 자신의 도메인 데이터에서도 마스킹 비율 트레이드오프가 같은 방향으로 나타나는지, 그리고 증류를 낮은 비율에만 걸었을 때 검색 성능이 유지되는지가 핵심 확인 지점이다.

한계와 전제도 분명하다. 사전학습 데이터는 웹, 코드, 학술, 지시 데이터로 구성된 웹 규모 자료이므로 학습 데이터의 편향을 물려받을 수 있고, 저자들은 다운스트림 사용자가 자기 사용 사례에 맞게 별도로 평가하길 권한다. 인코더의 다운스트림 상한은 관찰된 현상이며 추가 토큰으로 올릴 수 없다는 것이 이 논문의 전제다. RTD 비교에서는 CoLA를 교란 변수로 보고 부록에서 따로 다룬다. 증류는 어텐션 점수나 Gram 행렬 기반 방식이 길이 L에 대해 O(L²) 메모리와 연산을 요구해 긴 컨텍스트에 쓸 수 없고 FlashAttention과도 호환되지 않아, 로짓 기반으로 제한했다. 모델 가중치와 평가 코드, 사전학습 데이터 혼합은 공개한다.