Controlling Distributional Bias in Multi-Round LLM Generation via KL-Optimized Fine-Tuning
This addresses distributional bias in LLMs for applications requiring statistical alignment, though it appears incremental as it builds on existing fine-tuning and optimization techniques.
The paper tackles the problem of controlling output distributions in multi-round LLM generation, showing that standard methods fail to reliably align with target distributions like gender, race, and sentiment statistics. Their proposed fine-tuning framework with a hybrid objective significantly outperforms baselines, achieving precise distributional control across six datasets.
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.