Smooth MMD 정렬을 통한 LLM의 수치 예측 능력 향상

대규모 언어 모델(LLMs)은 강력한 일반 능력을 갖추고 있음에도 불구하고, 출력이 수치적으로 정밀해야 하는 경우에는 종종 신뢰하기 어려운 모습을 보입니다. 주요 원인은 학습 목적 함수(training objective)에 있습니다. 표준 교차 엔트로피(cross-entropy)는 숫자 토큰을 구조화되지 않은 범주로 취급하며, 해당 값들의 메트릭 구조(metric structure)를 무시합니다. 우리는 숫자 토큰에 대한 값-거리 커널(value-distance kernels)과 그래프 기반의 매끄러움(smoothness)을 결합하여 기존의 MMD를 확장한 Smooth Maximum Mean Discrepancy (SMMD)를 통해 이러한 불일치 문제를 해결합니다. 숫자 하위 어휘집(numeric sub-vocabulary)에 대해 정의된 이 커널을 사용하여, SMMD는 커널 매칭(kernel matching)을 통해 예측된 숫자 분포를 목표값에 정렬하고, 유도된 커널 그래프(kernel graph) 상에서 예측-목표 잔차(prediction-target residual)를 매끄럽게 만들어 국소적 일관성(local consistency)을 장려합니다. 우리는 여러 오픈 웨이트(open-weight) LLM 및 VLM 백본을 대상으로 수학적 추론(mathematical reasoning), 산술 계산(arithmetic calculation), 시계 시간 인식(clock-time recognition), 차트 질의응답(chart question answering) 등 네 가지 수치 목표 태스크에서 SMMD를 평가합니다. SMMD는 교차 엔트로피(cross-entropy) 및 최근의 수치 목표 손실 함수(numeric-target losses)보다 일관되게 높은 정확도를 향상시킵니다. 분석 결과, MMD와 매끄러움(smoothness) 사이의 상호 보완적인 효과가 확인되었으며 거리 기반 커널 설계의 중요성이 강조되었습니다. 코드는 https://github.com/Zuozhuo/smmd-loss 에서 확인할 수 있습니다.