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

arXiv: https://arxiv.org/abs/2604.05756 · HTML: https://mlsanity.com/paper/2604.05756/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 52 references that are in the map.
- First posted on arXiv 7 Apr 2026.

## Builds on (most cited first, 21 of 21)

- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [LaMP: When Large Language Models Meet Personalization](https://mlsanity.com/paper/2304.11406/) (2023)
- [RLPrompt: Optimizing Discrete Text Prompts with Reinforcement Learning](https://mlsanity.com/paper/2205.12548/) (2022)
- [Bias in Bios: A Case Study of Semantic Representation Bias in a High-Stakes Setting](https://mlsanity.com/paper/1901.09451/) (2019)
- [KTO: Model Alignment as Prospect Theoretic Optimization](https://mlsanity.com/paper/2402.01306/) (2024)
- [Queens are Powerful too: Mitigating Gender Bias in Dialogue Generation](https://mlsanity.com/paper/1911.03842/) (2019)
- [How well can Text-to-Image Generative Models understand Ethical Natural Language Interventions?](https://mlsanity.com/paper/2210.15230/) (2022)
- [Bias Testing and Mitigation in LLM-based Code Generation](https://mlsanity.com/paper/2309.14345/) (2023)
- [No Preference Left Behind: Group Distributional Preference Optimization](https://mlsanity.com/paper/2412.20299/) (2024)
- [Prompt-Based Length Controlled Generation with Multiple Control Types](https://mlsanity.com/paper/2406.10278/) (2024)
- [Controlled Text Generation for Large Language Model with Dynamic Attribute Graphs](https://mlsanity.com/paper/2402.11218/) (2024)
- [Biased Tales: Cultural and Topic Bias in Generating Children's Stories](https://mlsanity.com/paper/2509.07908/) (2025)
- [Enhancing Semantic Consistency of Large Language Models through Model Editing: An Interpretability-Oriented Approach](https://mlsanity.com/paper/2501.11041/) (2025)
- [Dynamic Multi-Reward Weighting for Multi-Style Controllable Generation](https://mlsanity.com/paper/2402.14146/) (2024)
- [Estimating LLM Consistency: A User Baseline vs Surrogate Metrics](https://mlsanity.com/paper/2505.23799/) (2025)
- [Prefix-Tuning Based Unsupervised Text Style Transfer](https://mlsanity.com/paper/2310.14599/) (2023)
- [Beyond Perception: Evaluating Abstract Visual Reasoning through Multi-Stage Task](https://mlsanity.com/paper/2505.21850/) (2025)
- [Establishing Trustworthiness: Rethinking Tasks and Model Evaluation](https://mlsanity.com/paper/2310.05442/) (2023)
- [LLM Bias Detection and Mitigation through the Lens of Desired Distributions](https://mlsanity.com/paper/2510.06354/) (2025)
- [Gender Encoding Patterns in Pretrained Language Model Representations](https://mlsanity.com/paper/2503.06734/) (2025)

## Built on by (most cited first, 2 of 2)

- [Who Should Be Generated? Justifying Demographic Targets in Open-Ended Generation](https://mlsanity.com/paper/2608.02551/) (2026)
- [MaD-RL: Matching Distributions for Calibrating LLMs with Reinforcement Learning](https://mlsanity.com/paper/2609.31644/) (2026)

## Abstract (arXiv)

While the real world is inherently stochastic, Large Language Models (LLMs) are predominantly evaluated on single-round inference against fixed ground truths. In this work, we shift the lens to distribution alignment: assessing whether LLMs, when prompted repeatedly, can generate outputs that adhere to a desired target distribution, e.g. reflecting real-world statistics or a uniform distribution. We formulate distribution alignment using the attributes of gender, race, and sentiment within occupational contexts. Our empirical analysis reveals that off-the-shelf LLMs and standard alignment techniques, including prompt engineering and Direct Preference Optimization, fail to reliably control output distributions. To bridge this gap, we propose a novel fine-tuning framework that couples Steering Token Calibration with Semantic Alignment. We introduce a hybrid objective function combining Kullback-Leibler divergence to anchor the probability mass of latent steering tokens and Kahneman-Tversky Optimization to bind these tokens to semantically consistent responses. Experiments across six diverse datasets demonstrate that our approach significantly outperforms baselines, achieving precise distributional control in attribute generation tasks.