고려대, 분류와 검색 성능을 동시에 잡은 UltraBERT 공개

고려대 DMIS 연구실이 영어 양방향 인코더 UltraBERT-150M/380M 및 임베딩 모델 UltraBERT-Embed-150M을 공개했다. 이 모델들은 32K 토큰 컨텍스트를 지원하며 Apache 2.0 라이선스로 배포되어 상업적 사용이 가능하다.

핵심 기술 및 구조

분류는 낮은 마스킹 비율, 검색은 높은 마스킹 비율에서 성능이 좋다는 발견에 기반하여 한 배치에 **5%**와 50% 마스킹 문장을 섞고 비율마다 별도의 LM 헤드를 사용하는 방식을 채택했다. 구조는 24층 고정이며 폭만 조절해 150M(은닉 768)과 380M(은닉 1024) 두 가지 크기로 제공된다. 어휘는 SmolLM2 토크나이저를 사용하며 국소 및 전역 어텐션을 교차 적용한다.

학습은 1.3B 교사 모델을 가지치기 후 직접 증류하는 Cascade Distillation 방식을 따랐다. 5% 마스킹 문장은 교사 로짓과의 KL divergence로 분류 헤드를 학습하고, 50% 마스킹 문장은 일반 MLM 손실로 검색 헤드를 학습했다. 최종 모델 학습에는 2T 토큰이 사용되었으며 마지막 100B 토큰에서 컨텍스트를 32K로 확장했다.

성능 및 벤치마크

UltraBERT-380M은 GLUE 91.0, BEIR 47.7을 기록하며 10개 벤치마크 중 9개에서 1위를 차지했다. 특히 코드와 수학 검색 벤치마크인 RAR-b에서 Code 47.0, Math 61.7로 2위 대비 각각 8.4점, 7.7점 높은 성능을 보였다. UltraBERT-150M은 분류에서 RoBERTa-large와 유사한 성능을, 검색에서는 Ettin-400M보다 우수한 성능을 달성했다.

32K 컨텍스트 성능 테스트인 LongEmbed passkey 과제에서는 256토큰부터 32K토큰 전 구간에서 **99%**의 정확도를 유지했다. 이는 ModernBERT나 Ettin이 16K/32K에서 급락하는 것과 비교된다. 임베딩 모델 UltraBERT-Embed-150M은 MTEB(eng) 67.24, MTEB(Code) 73.30을 기록하며 200M 이하 밀집 임베딩 모델 중 1위를 차지했다.

사용 시 주의사항

모델은 영어 및 코드 전용으로 한국어는 지원되지 않는다. 백본 모델(150M/380M)은 미세조정이 필수이며, 바로 사용 가능한 임베딩 모델은 UltraBERT-Embed-150M을 사용해야 한다. 학습률 설정 시 다른 인코더보다 최대 10배 낮게 설정해야 하며, 검색 미세조정의 경우 150M은 5e-5, 380M은 3e-5가 최적값이다. 또한 trust_remote_code=True 옵션이 필요하므로 운영 환경에서는 revision을 고정하는 것이 권장된다.