💬 NLP

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

This paper proposes a novel fine-tuning framework combining Steering Token Calibration with a hybrid KL-divergence and Kahneman-Tversky optimization objective to effectively control and align the output distributions of Large Language Models across multiple rounds, addressing the limitations of existing methods in generating statistically representative content regarding attributes like gender, race, and sentiment.

Original authors: Yanbei Jiang, Amr Keleg, Ryandito Diandaru, Jey Han Lau, Lea Frermann, Biaoyan Fang, Fajri Koto

Published 2026-04-08

📖 4 min read☕ Coffee break read

CC0 1.0

Original authors: Yanbei Jiang, Amr Keleg, Ryandito Diandaru, Jey Han Lau, Lea Frermann, Biaoyan Fang, Fajri Koto

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very talented, but slightly stubborn, storyteller named "LLM." You ask this storyteller to write 100 short stories about a specific job, like a "Nurse in the UK."

The Problem: The Stubborn Storyteller
In the real world, nursing is a mix of men and women. But if you ask your LLM storyteller to write 100 stories independently, it might accidentally write 90 stories about female nurses and only 10 about male nurses. It's not because the model is "evil"; it's because it has learned from the internet that "nurse" usually equals "woman."

If you ask it again for 100 stories about "Construction Workers," it might write 99 stories about men and 1 about a woman, even though you wanted a balanced mix.

The problem is that current AI models are great at telling one good story, but they are terrible at managing the big picture when asked to tell many stories. They tend to get stuck in a "rut," repeating the same stereotypes over and over, rather than reflecting the true diversity of the real world or a specific goal you set for them.

The Solution: The "Traffic Cop" and the "Editor"
The researchers in this paper came up with a clever two-part training method to fix this. Think of it as giving the AI a new set of tools: a Traffic Cop and a Strict Editor.

1. The Traffic Cop (Steering Tokens & KL Loss)

First, they teach the AI to use a special "secret code" or Traffic Cop token before it starts writing the story.

How it works: Before writing a story about a nurse, the AI must first pick a token. Let's say it picks Token "A" for "Male" or Token "B" for "Female."
The Goal: The researchers train the AI so that if you want 70% male nurses and 30% female nurses, the Traffic Cop picks Token "A" 70 times out of 100 and Token "B" 30 times out of 100.
The Analogy: Imagine a casino dealer. If you want the cards to be dealt in a specific ratio (e.g., 70% red, 30% black), the dealer (the Traffic Cop) is trained to shuffle and deal exactly that way, regardless of what the players usually expect.

2. The Strict Editor (Semantic Alignment & KTO)

Here is the tricky part. Just because the AI picked the "Male" token doesn't mean the story it writes afterwards will actually be about a man. It might pick the token by accident and then write a story about a woman anyway.

The Fix: They add a "Strict Editor" (using a technique called KTO). This editor checks: "Did you pick the 'Male' token? Great! Now, does the story you wrote actually describe a man?"
The Goal: If the AI picks "Male" but writes about a woman, the Editor gives it a big "F" (a penalty). If it picks "Male" and writes about a man, it gets an "A" (a reward).
The Analogy: Think of a chef who is told to make a "Spicy" dish. The Traffic Cop ensures they pick the "Spicy" label 50% of the time. The Strict Editor tastes the food to make sure it actually tastes spicy. If the chef picks the label but serves a bland soup, the Editor sends it back to the kitchen.

The Result: A Perfectly Balanced Menu
By combining the Traffic Cop (who controls the frequency of the choices) and the Strict Editor (who ensures the content matches the choice), the researchers created an AI that can:

Mirror Reality: If you ask for real-world statistics (e.g., "Write stories about construction workers based on US census data"), the AI will generate a mix that matches the real numbers perfectly.
Enforce Fairness: If you want a perfectly balanced mix (50/50) to avoid bias, the AI can do that too, even if the real world isn't balanced.

Why This Matters
Previously, if you wanted an AI to be fair or statistically accurate over many generations, you had to hope for the best or try to trick it with complex prompts (which usually failed). This new method is like giving the AI a dial that you can turn to set the exact percentage of men, women, or different sentiments you want to see in its output, without breaking its ability to tell a good story.

It's the difference between asking a chaotic crowd to "be diverse" (which usually results in a mess) and hiring a professional event planner who knows exactly how many people of each type to invite to make the party perfect.

1. Problem Statement

While Large Language Models (LLMs) are inherently probabilistic, standard evaluation methods typically focus on single-round inference against fixed ground truths. This paper identifies a critical gap: Distributional Bias in multi-round generation.

The Issue: When an LLM is prompted repeatedly (e.g., 100 times) with the same instruction (e.g., "Write a story about a nurse"), the aggregate output distribution often deviates significantly from the intended target distribution (e.g., real-world statistics or a uniform distribution).
The Consequence: Existing models and standard alignment techniques (like Prompt Engineering or Direct Preference Optimization) fail to reliably control these aggregate distributions. They tend to "mode-seek," converging on a single stereotypical response rather than maintaining a diverse distribution that reflects reality or user specifications.
Goal: To develop a framework that allows LLMs to generate outputs over multiple independent trials that strictly adhere to a user-defined target distribution (e.g., 96% male/4% female for construction workers, or 50/50 for educational content) while maintaining semantic quality.

2. Methodology

The authors propose a novel fine-tuning framework that couples Steering Token Calibration with Semantic Alignment using a hybrid objective function.

A. Core Components

Steering Tokens ( $S$ ):
- A set of special tokens (e.g., $\tau_1, \tau_2$ ) is added to the model's vocabulary, each uniquely corresponding to a specific attribute (e.g., "Male," "Female").
- These tokens act as precursors to the response, explicitly conditioning the generation on a specific attribute.
Two-Stage Optimization Strategy:
The training objective combines two distinct loss functions:
- KL Calibration Loss ( $L_{KL}$ ):
  - Purpose: To force the model's probability of generating a specific steering token to match the target distribution $P^*$ .
  - Mechanism: Minimizes the Kullback-Leibler (KL) divergence between the target distribution and the model's predicted probability of the steering tokens immediately following the prompt.
  - Effect: Acts as a "macro-controller" ensuring the statistical intent matches the target.
- Semantic Alignment via KTO ( $L_{KTO}$ ):
  - Purpose: To ensure the text generated after the steering token is semantically consistent with that token.
  - Mechanism: Uses Kahneman-Tversky Optimization (KTO). Unlike standard DPO which compares pairs, KTO defines a value function for individual examples based on whether they are desirable or undesirable.
  - Implementation:
    - Desirable ( $y_w$ ): Sequence $[\tau_{target}; \text{consistent response}]$ .
    - Undesirable ( $y_l$ ): Sequence $[\tau_{reject}; \text{inconsistent response}]$ .
  - Effect: Acts as a "micro-controller" ensuring the generated text logically follows the chosen steering token.
- Total Objective: $L_{total} = L_{KTO} + L_{KL}$

B. Data Construction

Base Dataset: Generated $N$ responses per instruction, explicitly controlled to match the target distribution $P^*$ .
Preference Pairs: Transformed into $(x, y_w, y_l)$ pairs. The accepted response ( $y_w$ ) prepends the correct steering token to the response. The rejected response ( $y_l$ ) prepends a negative steering token (sampled uniformly from remaining attributes) to a generated response, providing a contrastive signal.

C. Inference Strategy

A two-step generation process is used:

Token Sampling: Sample a steering token $\hat{\tau}$ from the model's probability distribution (which has been calibrated to match $P^*$ ).
Response Generation: Append $\hat{\tau}$ to the prompt and generate the rest of the response autoregressively.

3. Key Contributions

Problem Formulation: The first work to systematically analyze and address distributional bias in multi-round LLM generation, distinguishing it from single-round consistency or standard debiasing.
Novel Framework: Introduction of a hybrid fine-tuning method combining KL Divergence (for distribution calibration) and KTO (for semantic consistency), which outperforms standard preference optimization.
Comprehensive Evaluation: Creation of distribution-aware preference datasets covering Gender, Race, and Sentiment across UK and US occupational contexts.
Empirical Validation: Demonstration that standard baselines (Zero-shot, Prompt Engineering, IFT, DPO) fail to control distributions, while the proposed method achieves precise alignment.

4. Experimental Results

The method was evaluated on four models (Qwen2.5-7B/1.5B, Llama-3.1-8B/3.2-1B) across six datasets (Gender/Race/Sentiment in UK/US).

Performance Metric: Mean Absolute Error (MAE) between the generated distribution and the target distribution. Lower is better.
Main Findings:
- Superiority: The proposed method consistently achieved the lowest MAE across all models and settings. For example, on Qwen2.5-7B, it reduced average MAE from 0.135 (IFT) to 0.085 (a 37% improvement).
- Baseline Failure:
  - Prompt Engineering: Failed to reliably shift distributions.
  - DPO: Performed catastrophically (often reaching MAE of 0.50 in uniform settings), confirming that standard preference optimization forces mode-seeking behavior, collapsing the distribution to a single attribute.
  - IFT: Improved over zero-shot but significantly underperformed compared to the proposed hybrid method.
- Robustness: The method remained stable across different decoding parameters (Temperature, Top-p, Top-k) and generalized to unseen occupations (trained on UK, tested on US).
- General Capabilities: Evaluation on MMLU, GSM8K, TruthfulQA, and IFEval showed negligible degradation in general model capabilities.
Ablation Study:
- Removing $L_{KL}$ caused the highest performance drop, proving that token calibration is foundational.
- Removing $L_{KTO}$ degraded performance, showing the necessity of semantic binding.
- Replacing KTO with DPO led to instability, reinforcing the need for KTO's value-based approach over pairwise margins.

5. Significance

Theoretical: This work bridges the gap between distributional statistics and generative modeling, showing that LLMs can be trained to respect aggregate statistical constraints without sacrificing individual sample quality.
Practical:
- Fairness & Representation: Enables applications where LLMs must reflect real-world demographics (e.g., medical diagnostics, policy simulation) or enforce equity (e.g., educational materials requiring balanced representation).
- Control: Provides a reliable mechanism for "steering" LLM outputs in multi-round scenarios, a common use case in simulation, data augmentation, and interactive agents.
Ethical: The authors highlight the dual-use nature of this technology: it can be used to mirror real-world data (potentially reinforcing biases) or to enforce user-specified equitable distributions (debiasing). The framework provides the control necessary to make these ethical choices programmatically.

In conclusion, the paper presents a robust, mathematically grounded solution to a previously unaddressed limitation in LLMs: the inability to control the statistical distribution of outputs over repeated generations. By decoupling distribution calibration from semantic generation via a hybrid loss function, the authors achieve precise, stable, and generalizable control over model behavior.