기술 요약: 확률적 보정은 언어 모델에서 학습 가능한 능력이다

문제 정의

언어 모델 (LM) 은 다양한 출력 생성, 선택 확률 균형 맞추기, 또는 특정 수학 분포에서 샘플링과 같이 사용자가 지정하는 무작위성이 필요한 시나리오에서 점점 더 많이 배포되고 있다. 그러나 표준 자기회귀형 언어 모델은 본질적으로 열악한 확률적 생성기이다. 이들이 유도하는 출력 분포는 종종 좁은 부분 집합으로 붕괴되고, 지지 제약 조건을 위반하며, 요청된 목표 분포와 정렬되지 않는다. 기존 프롬프트 우회 기법 (예: 추론이나 시드 주입) 은 이러한 문제를 부분적으로 완화할 수 있지만, 추론 시 비용, 프롬프트 민감도, 계산 오버헤드를 도입할 뿐 모델의 샘플링 행동을 근본적으로 변경하지는 못한다.

이 논문이 다루는 핵심 질문은 확률적 보정—즉, 알려진 목표 분포를 반영하는 샘플을 생성하는 언어 모델의 능력—이 추론 시 휴리스틱에 의존하기보다는 파인튜닝을 통해 직접 학습될 수 있는지 여부이다.

방법론: 보정 파인튜닝

저자들은 수학 분포 (예: "[분포] 분포의 [파라미터] 를 가진 [분포] 에서 정확히 하나의 무작위 수를 생성하라") 에서 샘플링을 명시적으로 요청하는 합성 프롬프트를 사용하여 언어 모델을 보정된 확률적 샘플러로 훈련하는 프레임워크를 제안한다.

고정된 베이스 모델에 저랭크 적응 (LoRA) 을 활용하는 두 가지 다른 파인튜닝 전략이 비교되었다.

1. 소프트 타겟 파인튜닝 (Soft-Target Fine-Tuning):

   • 메커니즘: 목표 분포는 유한한 표준 출력 공간으로 이산화된다. 모든 유효한 수치 출력으로부터 접두사 트라이 (prefix trie) 가 구성되며, 각 노드는 토큰 접두사를 나타낸다.
   • 감독: 주어진 접두사에 대해, 트라이 내의 모든 유효한 연속에 대한 나머지 확률 질량을 정규화하여 목표 다음 토큰 분포를 계산한다.
   • 목적: 모델은 방문한 모든 접두사에서 예측된 다음 토큰 분포와 트라이에서 유도된 목표 분포 간의 커널-라이블 (KL) 발산을 최소화한다. 이는 밀집된 접두사 수준의 감독을 제공한다.

2. 하드 타겟 파인튜닝 (Hard-Target Fine-Tuning):

   • 메커니즘: 모델은 목표 분포에서 직접 샘플링된 완성본으로 훈련된다.
   • 감독: 각 프롬프트에 대해 목표 분포에서 수치 답이 샘플링되어 토큰화되고, 실제 완성본 (ground-truth completion) 으로 사용된다.
   • 목적: 모델은 샘플링된 완성본에 대한 표준 마스킹 자기회귀 교차 엔트로피를 최소화한다. 이는 희소하고 경로 수준의 감독을 제공하며, 동등한 커버리지를 달성하기 위해 더 많은 훈련 단계 (프롬프트당 더 많은 샘플) 를 필요로 한다.

실험 설정:

• 모델: 0.6B 에서 27B 파라미터까지 범위를 가진 4 개 계열 (Qwen3, Gemma-3, Llama-3.2, GPT-OSS) 의 12 개 모델.
• 데이터: 30 개 분포 계열 (24 개는 훈련용, 6 개는 분포 외 (OOD) 테스트용 홀드아웃).
• 평가:
  • 구조화된 샘플링: 목표 분포에 대한 워asserstein-1 거리 (W1) 및 로그 레벨 KL 발산을 통해 측정.
  • 전이 벤치마크: 오픈 엔디드 무작위 생성 (다양성 지지), MCQ 답안 위치 균형, 그리고 NoveltyBench (의미적 다양성).
  • 유지: TinyBenchmarks(MMLU, GSM8K 등) 에 대한 능력 유지 및 PALOMA 에 대한 언어 모델 적합도.

주요 기여

1. 보정의 학습 가능성: 이 논문은 확률적 보정이 학습 가능한 능력임을 입증한다. 소프트 타겟과 하드 타겟 파인튜닝 모두 홀드아웃 분포 계열과 보지 못한 파라미터 설정에서 구조화된 샘플링 충실도를 크게 향상시킨다.
2. 전략 비교 분석:
   • 하드 타겟: 일반적으로 도메인 내 구조화된 수치 샘플링 벤치마크, 특히 보지 못한 파라미터에 대해 우수한 성과를 거둔다. 정확한 수치 충실도에서 탁월하다.
   • 소프트 타겟: 수치 작업에서 경쟁력을 가지지만, 더 넓은 확률적 생성 설정으로의 전이에서 더 우수한 성능을 보인다. 오픈 엔디드 생성에서 일관되게 확률적 지지를 확장하고, MCQ 에서 답안 위치 균형을 개선하며, NoveltyBench 에서 더 높은 유용성을 제공한다.
3. 트레이드오프와 비용: 연구는 보정 향상이 종종 트레이드오프를 수반함을 강조한다. 구체적으로 산술 추론 (GSM8K) 은 경향성이 감소하는 경향이 있으며, 비용은 모델 크기와 파인튜닝 변형에 따라 달라진다. 그러나 언어 모델 적합도 (PALOMA 퍼플렉시티) 는 종종 유지되거나 향상되어, 모델이 단순히 토큰 확률을 평탄화하는 것이 아니라 특정 확률적 행동을 학습함을 시사한다.

결과

• 구조화된 샘플링: 두 방법 모두 베이스 체크포인트에 비해 정규화된 W1 및 로그 KL 을 크게 감소시켰다. 하드 타겟 파인튜닝은 일반적으로 수치 벤치마크에서 가장 강력한 성능을 발휘하며, 많은 경우 오류를 10 배 감소시켰다.
• 오픈 엔디드 생성: 소프트 타겟 파인튜닝은 모든 모델에서 상위 90% 지지 크기 (확률 질량의 90% 를 커버하는 데 필요한 토큰 수) 를 일관되게 증가시켰으며, 종종 1~2 차수만큼 증가시켰다. 이는 도시, 동물, 이름과 같은 카테고리에서 더 다양하고 덜 집중된 출력으로 이어진다.
• MCQ 및 Novelty: 소프트 타겟 구성은 일반적으로 MCQ 생성에서 더 나은 답안 위치 균형을 제공하고 NoveltyBench 에서 더 높은 유용성 점수를 얻어, 학습된 확률적 행동이 암시적 자연어 제약 조건으로 전이됨을 나타낸다.
• 유지: 일반 능력에 대한 영향은 혼합적이다. 일부 중대형 모델은 TinyBenchmarks 에서 집계 점수를 유지하거나 향상시키는 반면, 소형 모델과 특정 작업 (특히 GSM8K) 은 저하를 보인다. 이는 추론 흔적 없이 짧은 직접 완성본을 목표로 하는 훈련 목표가 추론이 많은 작업과 충돌할 수 있음을 시사한다.

중요성과 주장

이 논문은 확률적 보정이 언어 모델의 고유한 한계가 아니라 명시적으로 훈련될 수 있는 능력이라고 주장한다. 유효한 출력 공간과 그 위의 확률 질량 할당을 모두 학습하는 문제로 문제를 재구성함으로써, 저자들은 언어 모델을 제어된 확률적 샘플러로 변환할 수 있음을 보여준다.

이 연구의 중요성은 추론 시 프롬프팅을 넘어 모델의 분포적 행동을 근본적으로 변경하는 훈련 시 솔루션으로 이동하는 데 있다. 저자들은 향상 효과가 선택적이며 모든 하류 작업을 보편적으로 개선하지는 않는다 (산술 추론이 주목할 만한 예외) 고 지적하지만, 특정 확률적 행동을 위해 모델을 조정할 수 있는 능력은 다양성, 공정성, 또는 엄격한 통계적 제약 준수가 필요한 응용 분야에 새로운 길을 연다. 이 연구는 향후 연구가 이러한 목표에 의해 학습된 메커니즘에 초점을 맞추고 능력 저하를 완화하기 위해 유지 인식 변형을 개발하는 데 집중해야 함을 시사한다.