숫자 순서를 인식하는 임베딩 모델 실험
Number-aware embeddings
핵심 요약
기존 임베딩 모델의 숫자 순서 인식 문제를 해결하기 위해 토크나이저와 학습 방식을 개선한 실험 공유.
- 숫자 인식 문제 — 기존 임베딩 모델은 숫자 크기 순서를 제대로 이해하지 못함.
- 학습 방식 개선 — 토크나이저와 예측 헤드를 수정하고 로그 크기 기반으로 인코딩하여 학습함.
- 성능 향상 — 커스텀 벤치마크에서 기존 모델 대비 숫자 정렬 정확도가 크게 개선됨.
- 실용적 활용 — 숫자 데이터가 많은 HTML 테이블 등에서 정량적 데이터 추출에 효과적임.
"a 500 hp car", "a 1,200 hp car", "a 73 hp car"의 임베딩 간 코사인 유사도를 보면, 임베딩 모델이 숫자 순서를 전혀 이해하지 못한다는 걸 알 수 있습니다. (Qwen과 ModernBERT 기반 임베딩으로 테스트함)
이는 MLM 사전 학습 단계에서 토크나이저와 로그 우도 손실이 순서(Order Of Magnitude) 예측보다 정확한 예측에 과도하게 보상을 주기 때문입니다.
이를 해결하기 위해 숫자에 대한 기본 토크나이저/예측 헤드를 재정의하고, 수정된 아키텍처를 3억 개의 토큰(그중 약 400만 개가 숫자)으로 MLM 파인튜닝했습니다.
결과는 성공적이었습니다. 숫자 패턴을 정규식으로 찾아 로그 크기로 표현하는 방식입니다. 각 숫자는 128개의 빈(bin)으로 부드럽게 인코딩(인접 빈 사이의 선형 보간)되며, 각 빈에 대한 임베딩 사전 항목이 생성됩니다. 디코딩도 비슷하게 128개의 출력 빈과 부드러운 CE 손실을 사용하는 분류-회귀 헤드를 사용했습니다.
MLM 사전 학습 모델을 임베딩 모델로 만드는 과정이 가장 흥미로웠습니다. JEPA를 시도했다가 실패해서 인코더/디코더 설정을 사용했는데, 잘 작동했습니다.
6시간의 H100 학습 후 최종 결과입니다. 커스텀 벤치마크(이 문장이 완전 레드 플래그인 거 알지만요)에서 문장 세 쌍을 올바르게 정렬할 확률이 59%였습니다. ModernBERT(평균 풀링)는 38%, BGE-base-v1.5(CLS)는 34%였습니다.
숫자 위주의 HTML 테이블에서 구조화된/정량적 데이터를 추출하는 데도 꽤 뛰어납니다.
(다소 학습이 부족한) 모델은 여기 있습니다: https://huggingface.co/edereynal/financial_bert
전체 엔지니어링 과정이 궁금하시면 블로그 포스트를 확인해 주세요. 기술적으로 꽤 밀도가 높지만 흥미로울 겁니다: https://www.eloidereynal.com/p/i-spent-1-year-trying-to-predict
