기술 요약: 이산 자기회귀 언어 모델에 거리 척도 가르치기

문제 제기

대규모 언어 모델 (LLM) 과 기타 이산 자기회귀 기반 모델은 어휘에서 토큰을 예측하며, 각 토큰을 원-핫 (one-hot) 타겟 분포에 대해 최적화된 이산 범주형 변수로 취급합니다. 이러한 공식화는 LLM 을 비전, 로봇공학, 다중모달 추론 등 다양한 도메인에 적응시키는 데 기여했지만, 토큰 간의 거리적 관계를 근본적으로 무시합니다.

정수 산술 추론, 객체 감지의 공간 좌표, 로봇 조작의 회전 각도, 또는 이미지 생성의 양자화된 임베딩과 같이 수치적 또는 거리적으로 구조화된 요소를 포함하는 작업에서는 토큰이 원-핫 타겟으로 축소될 때 고유한 거리 구조가 손실됩니다. 기존 훈련 방식은 "4"라는 예측이 "2"라는 예측보다 "5"라는 정답에 의미적으로 더 가깝다는 사실을 간과합니다. 이러한 한계는 거리가 유의미한 도메인에서 데이터 효율성과 성능을 저해합니다.

방법론: DIST2Loss

저자들은 이산 자기회귀 모델을 위한 거리 인식 목적 함수인 DIST2Loss(DIScreTized DISTance Loss) 를 제시합니다. 핵심 아이디어는 표준 원-핫 타겟 분포를 사전 정의된 토큰 거리에서 유도된 보상 가중 분포로 대체하는 것입니다.

핵심 공식

어휘 부분집합 $V_d$와 거리 공간 $(X, d)$가 주어졌을 때, DIST2Loss 는 정답 토큰 $x$가 주어졌을 때 시간 단계 $t$에서의 토큰 $v$에 대한 타겟 분포 $p_d(v|x, t)$를 구성합니다. 이 분포는 이산화된 지수족을 사용하여 정의됩니다:

$$p_d(v|x, t) = \frac{\exp\left(-\frac{d(v, x, t)}{\tau}\right)}{\sum_{v' \in V_d} \exp\left(-\frac{d(v', x, t)}{\tau}\right)}$$

여기서:

• $d(v, x, t)$는 후보 토큰 $v$와 정답 $x$ 사이의 거리입니다 (예: 좌표의 유클리드 거리, 회전 각의 각도 거리, VQ 코드의 임베딩 거리).
• $\tau$는 분포의 부드러움을 조절하는 온도 하이퍼파라미터입니다.

최종 훈련 목적 함수는 표준 교차 엔트로피 손실 ($L_{CE}$) 과 거리 기반 정규화를 결합합니다:
$$L = L_{CE} + \alpha L_{dist}$$
여기서 $L_{dist}$는 모델이 예측한 분포와 거리 인식 타겟 $p_d$ 사이의 KL 발산이며, $\alpha$는 두 항의 균형을 맞춥니다.

이론적 해석

DIST2Loss 는 **엔트로피 정규화 정책 최적화의 폐쇄형 해 (closed-form solution)**로 해석됩니다. 강화 학습 (RL) 에서 엔트로피 정규화는 정책이 너무 일찍 수축하는 것을 방지하여 탐색을 장려합니다. 엔트로피 정규화와 함께 기대 보상을 최대화하는 최적 정책은 $\exp(R(a)/\tau)$에 비례하는 것으로 알려져 있습니다. DIST2Loss 는 후보 토큰을 행동으로, 음의 거리를 보상 신호로 간주하여 이를 구체화합니다. 샘플링, 롤아웃, 반복적 업데이트가 필요한 기존 RL 과 달리, DIST2Loss 는 안정적이고 분산이 없는 분석적 타겟 분포를 제공하여 정책 그래디언트 추정기와 관련된 불안정성을 피합니다.

고차원 및 멀티-토큰 시퀀스 처리

이 프레임워크는 임베딩에 대한 거리 척도를 정의함으로써 (예: 코사인 유사도) 고차원 연속 벡터 (예: VQ 임베딩) 를 처리할 만큼 유연합니다. 멀티-토큰 구조 요소 (예: 다중 자리 숫자) 의 경우, 저자들은 손실을 직접 적용할 때 어떤 토큰이 오류를 유발했는지 모델이 구분하지 못하는 "신용 할당 문제 (credit assignment problem)"에 직면한다고 지적합니다. 이를 완화하기 위해 대조적 타겟 증강(음의 이웃 샘플링) 과 자리값 가중치(위치적 중요도에 따른 손실 가중) 와 같은 실용적인 우회책을 제안하지만, 주요 실험은 보상이 토큰 수준에서 자연스럽게 분해되는 작업에 초점을 맞춥니다.

주요 기여

1. DIST2Loss 프레임워크: 추가 데이터나 상당한 계산 오버헤드 없이 이산 자기회귀 훈련에 거리 구조를 통합하는 플러그 앤 플레이 목적 함수.
2. 이론적 연결: 거리 인식 감독과 엔트로피 정규화 RL 간의 연결을 확립하여, 토큰별 보상이 알려진 작업에 대한 강화 학습의 안정적 대안을 제시.
3. 일반화: 토큰이 유의미한 거리를 허용한다면 1 차원 스칼라 회귀부터 고차원 이미지 생성에 이르기까지 다양한 도메인에 적용 가능함을 입증.

실험 결과

저자들은 DIST2Loss 를 다섯 가지 서로 다른 도메인에서 평가하여 데이터 효율성과 하류 성능에서 일관된 개선을 보였습니다:

1. 합성 함수 회귀 (토이 작업): 1~10 개의 샘플만으로도 선형 회귀를 학습하는 메타 학습 설정에서, DIST2Loss 는 평균 절대 오차 (MAE) 와 제곱근 평균 제곱 오차 (RMSE) 측면에서 표준 SFT 및 어휘 제약 기반선보다 월등히 우수한 성능을 보이며 우수한 구조적 일반화를 입증했습니다.
2. 시각적 그라운딩: 객체 감지 사전 훈련 없이 RefCOCO/RefCOCO+/RefCOCOg 에서 Phi3V 를 파인튜닝한 결과, DIST2Loss 는 대규모 감지 데이터셋으로 사전 훈련된 최첨단 모델과Comparable 한 성능을 달성했습니다. 특히 두 모델 모두 잘못된 객체를 예측한 "어려운 경우"에서도 DIST2Loss 예측은 정답에 기하학적으로 더 가까웠습니다 (높은 IoU).
3. 로봇 조작: VIMABench 에서 DIST2Loss 는 로봇 행동 생성 (관절 각도 및 좌표) 의 성공률을 향상시켰습니다. 이러한 이득은 데이터가 부족한 환경 (1K 샘플) 에서 가장 두드러졌으며, 데이터가 제한될 때 거리 사전 지식의 가치를 강조했습니다.
4. 생성적 보상 모델링: LLM 정렬 (RLHF) 에 적용된 DIST2Loss 는 RewardBench 와 MT-Bench 에서 보상 모델링 성능을 향상시켰습니다. DIST2Loss 를 사용한 생성적 접근 방식은 표준 이진 분류기 및 SFT 기반선보다 우월했으며, 특히 "Chat Hard" 및 "Safety" 도메인에서 두드러졌습니다.
5. 이미지 생성: 벡터 양자화 (VQ) 표현을 사용하여 DIST2Loss 는 자기회귀 모델 (LlamaGen) 에서 이미지 표현 학습을 강화하여 다양한 모델 크기와 훈련 에포크에 걸쳐 SFT 기반선보다 낮은 FID 와 높은 IS 점수를 달성했습니다.

중요성 및 주장

본 논문은 거리 인식 감독이 이산 자기회귀 모델을 위한 원-핫 감독의 간단하고 일반적인 대안이라고 주장합니다. 그 중요성은 다음과 같은 능력에 있습니다:

• 고유 구조 활용: 표준 교차 엔트로피가 무시하는 좌표, 숫자, 임베딩과 같은 많은 비텍스트 작업에 내재된 거리적 관계를 활용.
• 데이터 효율성 향상: 손실 함수에 귀납적 편향을 직접 인코딩하여 더 적은 훈련 샘플로 더 나은 성능 달성.
• 훈련 안정화: 샘플링 및 정책 그래디언트의 불안정성 문제를 피하면서 보상 정렬을 위한 폐쇄형 안정적 대안 제공.

저자들은 이 방법이 토큰이 유의미한 거리를 허용할 때 (예: 수치 값, 좌표, 양자화된 임베딩) 가장 효과적이라고 강조합니다. 고유한 기하학이 없는 도메인에서는 이 방법이 자연스럽게 표준 원-핫 감독으로 축소됩니다. 이 연구는 광범위한 다중모달 및 추론 작업에 대해 거리 척도를 통합하는 것이 이산 기반 모델 훈련의 중요하지만 과거에 미활용되었던 핵심 요소임을 시사합니다.