Controlling Distributional Bias in Multi-Round LLM Generation via KL-Optimized Fine-Tuning
- URL: http://arxiv.org/abs/2604.05756v1
- Date: Tue, 07 Apr 2026 11:58:16 GMT
- Title: Controlling Distributional Bias in Multi-Round LLM Generation via KL-Optimized Fine-Tuning
- Abstract summary: We assess whether Large Language Models (LLMs) can generate outputs that adhere to a desired target distribution.<n>We formulate distribution alignment using the attributes of gender, race, and sentiment within occupational contexts.<n>We propose a novel fine-tuning framework that couples Steering Token with Semantic Alignment.
- Score: 37.39836347901149
- License: http://creativecommons.org/publicdomain/zero/1.0/
- Abstract: 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.
Related papers
- Unbiased Open World Regularization for Fair Self-Supervised Learning [11.74837809839014]
Self-supervised learning models are susceptible to learning spurious biases in the dataset.<n>We propose Unbiased Open World Regularization (UOWReg), an encoder-only framework.<n>We empirically validate this framework across both Gaussian and spherical latent spaces.
arXiv Detail & Related papers (2026-07-24T09:51:22Z) - DISA: Offline Importance Sampling for Distribution-Matching LLM-RL [56.9445657766829]
We introduce DISA, which moves this calibration problem outside the RL loop.<n> DISA draws proposal trajectories offline, estimates the partition function via importance sampling, and freezes the resulting partition-function estimate.<n>On two open-weight backbones across six math and three code benchmarks, DISA matches or exceeds the online-coupled distribution-matching baseline Flow.
arXiv Detail & Related papers (2026-05-17T07:14:44Z) - Evidence-based Distributional Alignment for Large Language Models [58.65469623911573]
We propose Evi-DA, an evidence-based alignment technique that improves the fidelity and robustness of LLM-based distribution estimation.<n>Given a target country, Evi-DA retrieves related World Values Survey items and their answer distributions, predicts a coarse Welzel value signature for each option, and infers the country-conditioned answer distribution in a structured format.
arXiv Detail & Related papers (2026-03-03T03:34:06Z) - SetPO: Set-Level Policy Optimization for Diversity-Preserving LLM Reasoning [50.93295951454092]
We introduce a set level diversity objective defined over sampled trajectories using kernelized similarity.<n>Our approach derives a leave-one-out marginal contribution for each sampled trajectory and integrates this objective as a plug-in advantage shaping term for policy optimization.<n>Experiments across a range of model scales demonstrate the effectiveness of our proposed algorithm, consistently outperforming strong baselines in both Pass@1 and Pass@K across various benchmarks.
arXiv Detail & Related papers (2026-02-01T07:13:20Z) - Decoupling the Effect of Chain-of-Thought Reasoning: A Human Label Variation Perspective [60.45433515408158]
We show that long Chain-of-Thought (CoT) serves as a decisive decision-maker for the top option but fails to function as a granular distribution calibrator for ambiguous tasks.<n>We observe a distinct "decoupled mechanism": while CoT improves distributional alignment, final accuracy is dictated by CoT content.
arXiv Detail & Related papers (2026-01-06T16:26:40Z) - Whatever Remains Must Be True: Filtering Drives Reasoning in LLMs, Shaping Diversity [13.211627219720796]
Reinforcement Learning (RL) has become the de facto standard for tuning LLMs to solve tasks involving reasoning.<n>We argue that RL implicitly optimize the "mode-seeking" or "zero-forcing" Reverse KL to a target distribution causing the model to concentrate mass on certain high-probability regions of the target while others.<n>In this work, we instead begin from an explicit target distribution, obtained by filtering out incorrect answers while neglecting the relative probabilities of correct ones.
arXiv Detail & Related papers (2025-12-05T18:56:40Z) - Judging with Confidence: Calibrating Autoraters to Preference Distributions [56.17041629492863]
We argue that a reliable autorater must learn to model the full distribution of preferences defined by a target population.<n>We present two learning methods tailored to different data conditions.<n>Our results show that finetuning autoraters with a distribution-matching objective leads to verbalized probability predictions that are better aligned with the target preference distribution.
arXiv Detail & Related papers (2025-09-30T20:36:41Z) - Probabilistic Token Alignment for Large Language Model Fusion [100.30692772017238]
Training large language models (LLMs) from scratch can yield models with unique functionalities and strengths, but it is costly and often leads to redundant capabilities.<n>A key challenge in existing model fusion is their dependence on manually predefined vocabulary alignment.<n>We propose the probabilistic token alignment method as a general and soft mapping for alignment, named as PTA-LLM.
arXiv Detail & Related papers (2025-09-21T23:18:24Z) - On Optimal Steering to Achieve Exact Fairness [29.589891801235083]
Empirically, our optimal steering techniques on both synthetic and real-world datasets improve fairness without diminishing utility.<n>We demonstrate affine steering of LLM representations to reduce bias in multi-class classification.
arXiv Detail & Related papers (2025-09-19T08:37:51Z) - Distribution-Dependent Rates for Multi-Distribution Learning [26.38831409926518]
Recent multi-distribution learning framework tackles this objective in a dynamic interaction with the environment.
We provide distribution-dependent guarantees in the MDL regime, that scale with suboptimality gaps and result in superior dependence on the sample size.
We devise an adaptive optimistic algorithm, LCB-DR, that showcases enhanced dependence on the gaps, mirroring the contrast between uniform and optimistic allocation in the multi-armed bandit literature.
arXiv Detail & Related papers (2023-12-20T15:50:16Z) - A Distributional Approach to Controlled Text Generation [3.279201607581627]
We propose a Distributional Approach to address Controlled Text Generation from pre-trained Language Models (LMs)
This view permits to define, in a single formal framework, "pointwise" and "distributional" constraints over the target LM.
We then perform experiments over distributional constraints, a unique feature of our approach, demonstrating its potential as a remedy to the problem of Bias in Language Models.
arXiv Detail & Related papers (2020-12-21T19:02:41Z) - Learning Invariant Representations and Risks for Semi-supervised Domain
Adaptation [109.73983088432364]
We propose the first method that aims to simultaneously learn invariant representations and risks under the setting of semi-supervised domain adaptation (Semi-DA)
We introduce the LIRR algorithm for jointly textbfLearning textbfInvariant textbfRepresentations and textbfRisks.
arXiv Detail & Related papers (2020-10-09T15:42:35Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.