Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
Abstract
Federated Learning (FL) offers a privacy-preserving pathway for aligning Large Language Models (LLMs); however, existing frameworks typically enforce a monolithic reward model, inevitably averaging out inherently conflicting user preferences (e.g., helpfulness vs. harmlessness). While Variational Preference Learning (VPL) offers a pathway to personalization, adapting it to decentralized settings presents a fundamental challenge: \textit{posterior collapse} driven by severe local data scarcity and heterogeneity. In this paper, we propose Federated Variational Preference Alignment with Gumbel-Softmax Prior (FedVPA-GP), a framework designed to disentangle diverse preferences without compromising privacy. To stabilize variational inference, we introduce a Federated Mixture Prior that enables clients to leverage the aggregate population distribution as a dynamic prior. Furthermore, we incorporate an Orthogonal Loss that explicitly enforces the separation of preference prototypes in the latent space. Experiments on the HH-RLHF dataset demonstrate that FedVPA-GP significantly outperforms monolithic baselines, successfully disentangling conflicting user intents and enabling dynamic preference switching.
Lay Summary
AI assistants like ChatGPT are typically tuned to be both helpful and harmless using human feedback. But people disagree on how to balance these goals — and pooling everyone's feedback into a central server raises serious privacy concerns. Methods that train on each user's own device avoid this, but they force the system toward a single "one-size-fits-all" assistant that smooths over genuine differences between users, satisfying no one in particular. We propose FedVPA-GP, a method that learns a distinct "preference style" for each user without ever moving their raw feedback off-device. Devices coordinate by exchanging only abstract summaries of what kinds of preferences they represent. Two ideas make this work: each device can borrow useful patterns from peers with similar tastes, and the system actively keeps different preference styles from blurring into each other. On a public dialogue dataset, our method produces responses that are simultaneously more helpful and more harmless than competing approaches, and adapts cleanly to new users it has never seen before.