Rethinking Personalized Reward Modeling for LLMs under Preference Heterogeneity via Group-Debiased Federated Learning
2026-08-03 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors study how to train large language models to understand user preferences when users disagree on what they like. They find that a single shared model trained across all users can quickly adapt to different preferences, even if those preferences conflict, due to the model’s flexible starting point. However, this works best when user groups are balanced; imbalanced groups cause problems. To fix this, the authors propose a new method called FedGD that adjusts how users are sampled during training to balance group influence, helping the model better personalize to everyone without needing to know user groups beforehand.
federated learningreward modelingpreference heterogeneityFedAvgmodel personalizationflatness of initializationgroup imbalanceclient samplinglatent preference groupslocal fine-tuning
Authors
Seongyoon Kim, Boryeong Cho, Jihwan Oh, Seokhyun Chung, Se-Young Yun
Abstract
Large language models are increasingly aligned to human preferences via reward modeling, but user preference data are sensitive and often cannot be centralized. Federated learning keeps such data local while learning a shared initial reward model, which is later personalized for each client through local fine-tuning. Because users often assign opposite labels to the same pair of responses, existing federated methods address preference heterogeneity by clustering similar clients and training one reward model per group, assuming that each group requires its own initialization. We show that this assumption is unnecessary. Under balanced preference groups, a single FedAvg model, despite starting at nearly random accuracy, surpasses reward models trained separately for each ground-truth group after only a few local optimization steps. We attribute this phenomenon to the flatness of the shared initialization: averaging across all clients learns richer shared representations that distinguish responses while canceling conflicting preference directions, leaving the model near a decision boundary that can be rapidly adapted. Group imbalance breaks this effect as the cancellation becomes asymmetric and leaves minority clients too far from the boundary to recover. Motivated by this observation, we propose FedGD (Federated Learning with Group Debiasing), which discovers latent preference groups during federated training and learns a single reward model using group-debiased client sampling. By counteracting the effect of group imbalance, FedGD learns an initialization that remains highly adaptable, enabling effective personalization without prior knowledge of the underlying groups.