Controlling Distributional Bias in Multi-Round LLM Generation via KL-Optimized Fine-Tuning

1. 문제 정의 (Problem Definition)

• 배경: 현실 세계는 본질적으로 확률적 (stochastic) 이지만, 현재 대규모 언어 모델 (LLM) 은 고정된 정답 (ground truth) 에 대한 단일 회차 추론 (single-round inference) 으로 주로 평가됩니다.
• 핵심 문제: 실제 응용 분야에서는 동일한 프롬프트를 반복적으로 입력할 때, LLM 의 출력 분포가 의도된 통계 (예: 현실 세계의 인구 통계 또는 균일 분포) 와 달라지는 **분포 편향 (Distributional Bias)**이 발생합니다.
  • 예시: "영국에서 간호사에 대한 이야기를 써라"라는 프롬프트를 100 회 독립적으로 입력했을 때, 모델이 성별 편향 (예: 여성 간호사만 계속 생성) 을 보일 수 있습니다.
• 한계: 기존 제로샷 (Zero-shot) 프롬프팅, 프롬프트 엔지니어링, 또는 표준 정렬 기법 (DPO 등) 은 이러한 다중 회차 생성에서의 분포를 안정적으로 제어하지 못합니다. 특히 DPO 와 같은 선호도 최적화 기법은 단일 최상의 응답에 수렴하려는 성향 (mode-seeking) 이 있어 전체 분포를 맞추는 데 실패합니다.

2. 제안 방법 (Methodology)

저자들은 **분포 정렬 (Distribution Alignment)**을 달성하기 위해 **스테어링 토큰 보정 (Steering Token Calibration)**과 **의미 정렬 (Semantic Alignment)**을 결합한 새로운 파인튜닝 프레임워크를 제안합니다.

2.1. 핵심 구성 요소

1. 스테어링 토큰 (Steering Tokens):

   • 속성 (예: 남성, 여성) 에 대응하는 특수 토큰 집합 $S = \{\tau_1, \tau_2, ...\}$을 모델 어휘에 추가합니다.
   • 이 토큰들은 응답 생성 전에 먼저 샘플링되어 모델의 "의도"를 제어합니다.

2. 하이브리드 목적 함수 (Hybrid Objective Function):

   • KL 발산 손실 (KL Loss, $L_{KL}$):
     • 목적: 모델이 생성하는 스테어링 토큰의 확률 분포가 목표 분포 $P^*$ (예: 현실 통계 또는 균일 분포) 와 일치하도록 보정합니다.
     • 역할: 거시적 (Macro) 인 분포 제어를 담당합니다.
     • 수식: $L_{KL}(\theta) = E_{x \sim D} [D_{KL} (P^*(\cdot|x) \parallel P_\theta(\cdot|x))]$
   • 카네만 - 트버스키 최적화 (KTO Loss, $L_{KTO}$):
     • 목적: 선택된 스테어링 토큰 뒤에 이어지는 텍스트 응답이 해당 속성과 의미적으로 일관되도록 합니다. (예: '남성' 토큰 뒤에는 남성 관련 내용이 나와야 함)
     • 역할: 미시적 (Micro) 인 의미 일관성을 보장하며, 기존 DPO 와 달리 개별 예제의 가치 (Value) 를 기반으로 학습합니다.
     • 수식: $L_{KTO}(\theta) = E [( \lambda_D - v(x, y_w)) + (\lambda_U - v(x, y_l))]$

3. 데이터 구성 (Data Construction):

   • Base Dataset: 목표 분포에 따라 속성별 응답 수를 조절하여 생성합니다.
   • Preference Pairs:
     • Accepted ($y_w$): 목표 속성에 맞는 스테어링 토큰 + 해당 응답.
     • Rejected ($y_l$): 나머지 속성 중 하나를 균일하게 샘플링하여 부정적 토큰 + 해당 응답.

4. 추론 전략 (Inference Strategy):

   • Step 1: 프롬프트 입력 후 스테어링 토큰의 확률을 계산하여 목표 분포에 맞게 토큰을 샘플링합니다.
   • Step 2: 샘플링된 토큰을 프롬프트에 추가한 후, 모델이 이를 조건으로 하여 응답을 생성합니다.

3. 주요 기여 (Key Contributions)

1. 다중 회차 생성에서의 편향 발견: 오프더셸 (off-the-shelf) LLM 들이 반복 생성 시 현실 통계나 균일 분포를 따르지 못하며, 기존 편향 완화 기법들 (프롬프팅, DPO 등) 이 이 문제를 해결하지 못함을 실증했습니다.
2. 새로운 데이터셋 구축: 영국 및 미국의 직업을 기반으로 성별, 인종, 감성 속성을 포함한 분포 인식 선호도 데이터셋을 구축했습니다.
3. 혁신적인 학습 프레임워크: KL 발산을 통한 분포 보정과 KTO 를 통한 의미 정렬을 결합한 하이브리드 파인튜닝 방법을 제안하여, LLM 의 출력 분포를 정밀하게 제어할 수 있음을 증명했습니다.

4. 실험 결과 (Results)

• 데이터셋: 미국/영국 직업 (39 개), 성별 (남/여), 인종, 감성 등 6 가지 시나리오.
• 평가 지표: 목표 분포와 생성된 분포 간의 편차를 측정하는 평균 절대 오차 (MAE).
• 성능:
  • 제안된 방법 ('Ours') 은 모든 모델 (Qwen, Llama) 및 모든 데이터셋에서 가장 낮은 MAE를 기록했습니다.
  • 비교:
    • Zero-shot / 프롬프트 엔지니어링: 분포 제어에 실패 (높은 MAE).
    • IFT (지도 학습): 어느 정도 개선되었으나 분포 정밀도 부족.
    • DPO: 분포 제어에 실패 (특히 균일 분포 설정에서 MAE 가 0.50 으로 극단적으로 높음, 단일 속성으로 수렴).
    • 제안 방법: Qwen2.5-7B 기준 평균 MAE 를 0.135(IFT) 에서 0.085 로 37% 개선.
• 추가 분석:
  • Ablation Study: KL Loss 제거 시 분포 정렬 실패, KTO 제거 시 의미 불일치 발생. 두 손실 함수의 시너지가 필수적임.
  • 일반화 능력: 학습하지 않은 직종 (Unseen Occupations) 에 대해서도 우수한 성능을 보임.
  • 일반 능력 유지: MMLU, GSM8K 등 일반 벤치마크에서 성능 저하가 거의 없음.

5. 의의 및 결론 (Significance)

• 실용적 가치: 의료 진단, 안전 관련 작업 등 현실 통계의 정확성이 중요한 분야에서 LLM 의 출력을 현실 세계 분포에 맞출 수 있게 합니다.
• 윤리적 통제: 반대로 편향을 제거하고 균일한 분포 (예: 교육용 텍스트에서의 성별 평등) 를 강제하여 사회적 편견을 완화하는 데 활용 가능합니다.
• 기술적 기여: LLM 의 확률적 생성 특성을 단순히 '불일치'로 보지 않고, **분포 제어 (Distributional Control)**라는 관점에서 접근하고 이를 정밀하게 제어할 수 있는 새로운 최적화 기법을 제시했습니다.

이 논문은 LLM 이 단일 답변의 정확성뿐만 아니라, 반복 생성 시 나타나는 전체적인 통계적 성향을 사용자가 의도한 대로 제어할 수 있는 길을 열었다는 점에서 중요한 의의를 가집니다.