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

Yanbei Jiang; Amr Keleg; Ryandito Diandaru; Jey Han Lau; Lea Frermann; Biaoyan Fang; Fajri Koto

cs.CLApr 7, 2026

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

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

Organizations: University of Melbourne · Mohamed bin Zayed University of Artificial Intelligence · Oracle

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.

Figures & tables

Figure 1: Distributional Bias in Multi-round LLM Generation

Figure 2: Mitigating Distributional Bias with KTO and KL Loss

Model	Method	Gender (UK)		Gender (US)		Race	Sentiment	Avg.
Model	Method	Real	Even	Real	Even	Even	Even	Avg.
Qwen2.5-7B-Instruct	Zero-shot	0.132	0.308	0.130	0.306	0.319	0.405	0.267
	PE-Explicit	0.209	0.331	0.231	0.409	0.240	0.444	0.311
	PE-Implicit	0.260	-	0.195	-	-	-	0.228
	IFT	0.144	0.080	0.247	0.051	0.123	0.162	0.135
	DPO	0.193	0.500	0.177	0.500	0.264	0.364	0.333

Table 1: Main results measured by MAE across four models and six datasets (lower is better). The best performance is highlighted in bold. “Real” denotes target distributions adhering to real-world statistics, while “Even” denotes uniform distribution targets. “Avg.” denotes average MAE over six datasets. Under the Even setting, PE-Explicit and PE-Implicit are identical; therefore, PE-Implicit is omitted in these cases. See Appendix A.2 for additional statistical analyses, including confidence intervals and standard deviations across occupations.

Figure 3: The representation of females in [0, 100] for the 14 considered occupations in the US. The first row represents the real-world statistics for each occupation.

Model & Method		Gen (UK)		Gen (US)		Senti	Avg.
		Real	Even	Real	Even	Even	Avg.
Qwen7B	Zero	0.26	0.34	0.26	0.38	0.42	0.33
	Exp	0.26	0.36	0.26	0.35	0.33	0.31
	Imp	0.44	-	0.43	-	-	0.44
	IFT	0.27	0.34	0.25	0.37	0.42	0.33
	DPO	0.23	0.31	0.26	0.36	0.42	0.32

Table 2: Story generation results measured by MAE across four models and five datasets. The best performance is highlighted in bold.

Model & Method		Gen (UK)		Gen (US)		Race	Senti
		Real	Even	Real	Even	Even	Even
Qwen7B	w/o KL	0.19	0.50	0.28	0.08	0.15	0.36
	w/o KTO	0.12	0.04	0.28	0.07	0.14	0.09
	swap DPO	0.18	0.26	0.19	0.50	0.12	0.09
	Ours	0.09	0.05	0.09	0.06	0.11	0.11
Qwen1.5B	w/o KL	0.34	0.24	0.27	0.10	0.14	0.07

Table 3: Ablation study on Qwen models measured by MAE. “w/o KL” and “w/o KTO” denote removing the Steering Token Calibration loss and Semantic Alignment loss, respectively. “swap DPO” denote swapping the KTO loss to DPO loss.

Figure 4: Sensitivity of Top - $p$ , Top - $k$ , and Temperature on Qwen2.5 - 1.5B with UK Gender Real dataset. Y - axis: MAE; shaded areas: std. dev. over 5 runs.

Model & Method		Gen (UK)		Gen (US)		Avg.
		Real	Even	Real	Even	Avg.
Qwen7B	Zero	0.25	0.11	0.30	0.26	0.23
	IFT	0.15	0.19	0.12	0.12	0.14
	Ours	0.09	0.14	0.05	0.06	0.08
Qwen1.5B	Zero	0.12	0.14	0.29	0.30	0.21
	IFT	0.13	0.24	0.06	0.04	0.12

Table 4: Internal alignment measured by MAE based on the Softmax probabilities of gender tokens (Male/Female). The results reflect the model’s intrinsic probability landscape prior to decoding strategies.

Panel of Figure 5 — Figure 5: Steering token versus generated Male gender probability using the data with real-world UK distribution. The dashed diagonal represents perfect alignment.

Model & Method		MMLU	GSM8K	TruthfulQA	IFEval
Qwen7B	Zero	26.91	81.35	59.36	67.62
	IFT	27.57	81.73	57.41	66.31
	DPO	27.57	81.65	58.38	66.31
	Ours	26.87	81.12	60.00	65.23
Qwen1.5B	Zero	28.81	62.32	54.00	42.93
	IFT	27.80	61.64	55.51	39.21

Table 5: Evaluation of model capabilities across general domains using the fine-tuned model from UK Real dataset. Metrics reported are Accuracy for MMLU, GSM8K, and IFEval, and BLEU-RT for TruthfulQA.

Method	UK Real	UK Even	US Real	US Even
Ours (0/1)	0.084	0.048	0.158	0.054
Token (man/woman)	0.088	0.041	0.168	0.047
Token (he/she)	0.078	0.049	0.136	0.050

Table 6: Comparison of methods on UK and US datasets under Real and Even settings using Qwen2.5-1.5B.

Method	Model	Real	Even
IFT	Qwen2.5-7B	0.158	0.189
	Qwen2.5-1.5B	0.187	0.154
	Llama-3.1-8B	0.109	0.101
	Llama-3.2-1B	0.147	0.085
DPO	Qwen2.5-7B	0.201	0.500
	Qwen2.5-1.5B	0.265	0.500

Table 7: Comparison of IFT, DPO, and our method on UK datasets under Real and Even settings across different models.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset Split	Attribute	Train	Val	Test
UK Occupations (Real)	Gender	2500	300	2500
UK Occupations (Even)	Gender	2500	300	2500
US Occupations (Real)	Gender	1400	300	1400
US Occupations (Even)	Gender	1400	300	1400
US Occupations (Even)	Race	1400	300	1400
US Occupations (Even)	Sentiment	1400	300	1400

Appendix

Table 8: Statistics of the six dataset splits used in our experiments. “Real” targets denote real-world census distributions, while “Even” targets denote uniform distributions.

Attribute	Task Setting	Representative Prompt Template
Gender	Attribute Generation	What is the gender of the [OCCUPATION]?
	Story Generation	Write a short story about a day in the life of a [OCCUPATION].
Race	Attribute Generation	What is the racial background of the [OCCUPATION]?
	Story Generation	Describe a scene featuring a [OCCUPATION] at work.
Sentiment	Attribute Generation	How would you describe the mood of the [OCCUPATION]?
	Story Generation	Tell a story about a [OCCUPATION] that reflects a specific emotional tone.

Appendix

Table 9: Representative prompt templates for the 12 experimental settings (covering 6 dataset splits

\times

2 tasks). ‘[OCCUPATION]’ is replaced by the specific role from our US or UK occupation lists. Actual inputs include paraphrased variations of these templates.

Model	Method	Gender (UK)		Gender (US)		Race	Sentiment
Model	Method	Real	Even	Real	Even	Even	Even
Qwen2.5-7B	Zero-shot	$0.132_{\pm 0.109}$	$0.308_{\pm 0.151}$	$0.130_{\pm 0.070}$	$0.306_{\pm 0.163}$	$0.319_{\pm 0.002}$	$0.405_{\pm 0.029}$
	PE-Explicit	$0.209_{\pm 0.133}$	$0.331_{\pm 0.156}$	$0.231_{\pm 0.118}$	$0.409_{\pm 0.126}$	$0.240_{\pm 0.000}$	$0.444_{\pm 0.000}$
	PE-Implicit	$0.260_{\pm 0.196}$	-	$0.195_{\pm 0.135}$	-	-	-
	IFT	$0.144_{\pm 0.098}$	$0.080_{\pm 0.052}$	$0.247_{\pm 0.144}$	$\mathbf{0.051}_{\pm 0.143}$	$0.123_{\pm 0.030}$	$0.162_{\pm 0.046}$
	DPO	$0.193_{\pm 0.139}$	$0.500_{\pm 0.000}$	$0.177_{\pm 0.141}$	$0.500_{\pm 0.000}$	$0.264_{\pm 0.022}$	$0.364_{\pm 0.036}$

Appendix

Table 10: Distribution alignment performance measured by MAE and its Standard Deviation (SD) across occupations (

Mean_{\pm SD}

). Lower SD indicates that the method’s debiasing capability is robust and consistent across different occupation types, rather than being effective only on specific roles.

Model	Method	Gender (UK)		Gender (US)		Race	Sentiment
Model	Method	Real	Even	Real	Even	Even	Even
Qwen2.5-7B	Zero-shot	$0.132_{\scriptscriptstyle{[0.12,0.15]}}$	$0.308_{\scriptscriptstyle{[0.29,031]}}$	$0.130_{\scriptscriptstyle{[0.13,0.15]}}$	$0.306_{\scriptscriptstyle{[0.29,0.32]}}$	$0.319_{\scriptscriptstyle{[0.31,0.32]}}$	$0.405_{\scriptscriptstyle{[0.40,0.42]}}$
	PE-Explicit	$0.209_{\scriptscriptstyle{[0.21,0.22]}}$	$0.331_{\scriptscriptstyle{[0.31,0.33]}}$	$0.231_{\scriptscriptstyle{[0.23,0.24]}}$	$0.409_{\scriptscriptstyle{[0.40,0.41]}}$	$0.240_{\scriptscriptstyle{[0.24,0.25]}}$	$0.444_{\scriptscriptstyle{[0.44,0.44]}}$
	PE-Implicit	$0.260_{\scriptscriptstyle{[0.25,0.27]}}$	-	$0.195_{\scriptscriptstyle{[0.19,0.20]}}$	-	-	-
	IFT	$0.144_{\scriptscriptstyle{[0.13,0.16]}}$	$0.080_{\scriptscriptstyle{[0.06,0.09]}}$	$0.247_{\scriptscriptstyle{[0.24,0.27]}}$	$\mathbf{0.051}_{\scriptscriptstyle{[0.04,0.08]}}$	$0.123_{\scriptscriptstyle{[0.12,0.14]}}$	$0.162_{\scriptscriptstyle{[0.15,0.18]}}$
	DPO	$0.193_{\scriptscriptstyle{[0.17,0.19]}}$	$0.500_{\scriptscriptstyle{[0.50,0.50]}}$	$0.177_{\scriptscriptstyle{[0.16,0.20]}}$	$0.500_{\scriptscriptstyle{[0.50,0.50]}}$	$0.264_{\scriptscriptstyle{[0.26,0.28]}}$	$0.364_{\scriptscriptstyle{[0.35,0.38]}}$

Appendix

Table 11: Distribution alignment performance with 95% Confidence Intervals (

Mean_{\pm CI}

). CIs are calculated via bootstrapping (

k=1000

) on the test set. The tight intervals and minimal overlap with baselines reinforce the statistical significance of our method’s performance.

Figure 6: Sensitivity analysis of Top- $p$ , Top- $k$ , and Temperature on the Qwen2.5-1.5B model and the UK Gender Even Datasets.

Figure 7: Sensitivity analysis of Top- $p$ , Top- $k$ , and Temperature on the Qwen2.5-1.5B model and the US Gender Real Datasets.

Figure 8: Sensitivity analysis of Top- $p$ , Top- $k$ , and Temperature on the Qwen2.5-1.5B model and the US Gender Even Datasets.

Figure 9: Sensitivity analysis of Top- $p$ , Top- $k$ , and Temperature on the Qwen2.5-1.5B model and the Race Datasets.

Figure 10: Sensitivity analysis of Top- $p$ , Top- $k$ , and Temperature on the Qwen2.5-1.5B model and the Sentiment Datasets.

Model & Method		Gen (UK)		Gen (US)		Avg.
		Real	Even	Real	Even	Avg.
Qwen7B	Zero	0.25	0.11	0.30	0.26	0.23
	IFT	0.15	0.19	0.12	0.12	0.14
	Ours	0.09	0.14	0.05	0.06	0.08
Qwen1.5B	Zero	0.12	0.14	0.29	0.30	0.21
	IFT	0.13	0.24	0.06	0.04	0.12

Appendix

Table 12: Internal alignment measured by MAE based on the Softmax probabilities of gender tokens (Male/Female). The results reflect the model’s intrinsic probability landscape prior to decoding strategies.

Panel of Figure 11 — Figure 11: The representation of females in [0, 100] for the 25 occupations in the UK, and the 14 considered occupations in the US. The first row represents the real-world statistics for each occupation.

LLM Alignment

Large language model alignment studies methods for making models follow human intentions, values, and safety constraints across contexts.

See all 611 papers

Evidence-based Distributional Alignment for Large Language Models
Mar 3, 2026Viet-Thanh Pham, Lizhen Qu, Zhuang Li +1LLM Alignment Cross-Cultural Language Model Evaluation
Department of Data Science & AI, Monash University, Australia · Royal Melbourne Institute of Technology, Australia
MaD-RL: Matching Distributions for Calibrating LLMs with Reinforcement Learning
Sep 12, 2026Sourabh Kulkarni, Ksheeraj Sai Vepuri, Basar Demir +8Distribution Matching RL for Language Models
Meta Superintelligence Labs
Pref-CTRL: Preference Driven LLM Alignment using Representation Editing
Apr 26, 2026Imranul Ashrafi, Inigo Jauregi Unanue, Massimo PiccardiLLM Alignment Representation Engineering
University of Technology Sydney, Australia

Mar 3, 2026J/K move · Enter open · S save

Evidence-based Distributional Alignment for Large Language Models

Viet-Thanh Pham, Lizhen Qu, Zhuang Li, Gholamreza Haffari

Department of Data Science & AI, Monash University, Australia · Royal Melbourne Institute of Technology, Australia

LLM Alignment Cross-Cultural Language Model Evaluation Cultural Bias in Language Models

Open paper AskPDF arXiv

Distributional alignment enables large language models (LLMs) to predict how a target population distributes its responses across answer options, rather than collapsing disagreement into a single consensus answer. However, existing LLM-based distribution prediction is often unstable and degrades under cultural and domain shift. Token score-based estimates can change with minor option wording or formatting, response sampling-based estimates are expensive and sensitive to prompts and decoding settings, and directly generated distributions are frequently miscalibrated. We propose Evi-DA, an evidence-based alignment technique that improves the fidelity and robustness of LLM-based distribution estimation under domain and cultural shift. Given a target country and a multiple-choice question, 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. We train the LLMs using a two-stage pipeline, where reinforcement learning optimizes survey-derived rewards that encourage accurate intermediate value predictions, faithful final distributions, well-formed structured outputs, and reduced cultural bias. Across in-domain and out-of-domain benchmarks and multiple open-source backbones, Evi-DA reduces Jensen-Shannon divergence between predicted and gold distributions relative to strong baselines, with average relative improvements of up to 44%.

Abstract

LLM Alignment

Explore similar work

Evidence-based Distributional Alignment for Large Language Models

MaD-RL: Matching Distributions for Calibrating LLMs with Reinforcement Learning

Pref-CTRL: Preference Driven LLM Alignment using Representation Editing