Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences Jabin Koo 1 Hoyoung Kim 2 Minwoo Jang 3 Jungseul Ok 3 1
arXiv:2605.30873v1 [cs.LG] 29 May 2026
Abstract
a critical bottleneck: high-quality preference data—often reflecting personal, cultural, and political nuances—resides on edge devices. Centralizing this data not only risks severe privacy violations, such as the extraction of sensitive training data (Carlini et al., 2021), but also faces challenges with regulations like GDPR (European Parliament and Council of the European Union, 2016).
Federated Learning (FL) offers a privacypreserving 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: posterior collapse driven by severe local data scarcity and heterogeneity. In this paper, we propose Federated Variational Preference Alignment with Gumbel-Softmax Prior (FedVPAGP), 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.
Federated Learning (FL) offers a privacy-preserving alternative (McMahan et al., 2017; Kairouz et al., 2021). Recent frameworks have adapted alignment to this decentralized setting, such as FedDPO (Ye et al., 2024) and FedBiscuit (Wu et al., 2024). Notably, FedBiscuit addresses computational constraints by training only a binary preference selector on the client side while keeping the base LLM frozen (Wu et al., 2024). However, despite these advancements, existing frameworks share a critical limitation: they enforce a monolithic reward model. Human values are inherently pluralistic and can be conflicting—for instance, preferences often diverge between helpfulness and harmlessness (Santurkar et al., 2023; Poddar et al., 2024). Aggregating these heterogeneous distributions into a single global model yields theoretical sub-optimality (Shirali et al., 2025). By aiming for a monolithic solution, these methods implicitly enforce a consensus that does not exist, resulting in a one-size-fits-all model that fails to satisfy distinct client needs. While Variational Preference Learning (VPL) (Poddar et al., 2024) offers a pathway to personalization by modeling user intent as a latent variable, adapting it to the federated setting presents a fundamental challenge driven by two intrinsic characteristics of FL: data heterogeneity and data scarcity. In centralized regimes, the model learns a dense preference manifold from pooled data, allowing it to distinguish subtle variations in user intent. In contrast, federated clients operate on highly heterogeneous distributions, where each client observes only a fragmented slice of global preferences, often restricted to a single mode like helpfulness or harmlessness. Compounding this, the severe scarcity of local samples causes the KL regularization term to dominate the reconstruction objective during variational inference. Lacking both the global context to position their preferences and sufficient data to support complex posterior estimation, the latent variable often degenerates to an uninformative prior. This phenomenon, known as posterior collapse, renders
1. Introduction Reinforcement Learning from Human Feedback (RLHF) has established itself as the standard paradigm for aligning Large Language Models (LLMs) with human intent (Christiano et al., 2017; Ziegler et al., 2019; Ouyang et al., 2022). However, its reliance on centralized data aggregation poses 1 Department of CSE, POSTECH, Pohang, Republic of Korea National AI Research Lab, Seoul, Republic of Korea 3 Graduate School of AI, POSTECH, Pohang, Republic of Korea. Correspondence to: Jungseul Ok <[email protected]>. 2
Proceedings of the 43 rd International Conference on Machine Learning, Seoul, South Korea. PMLR 306, 2026. Copyright 2026 by the author(s).
1
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
Federated Variational Preference Alignment with Gumbel-Softmax Prior (FedVPA-GP) (a) Federated Variational Binary Selector Training Helpful Client 1
𝐷1
Preference
(b) Local Training Loss Gumbel-softmax Sampling
Harmless Client 2
𝐷2
Gumbel 0,1
Preference
Local Training
Variational Encoder: 𝑞𝜙1 𝑧 ∣ 𝐷1
Variational Encoder: 𝑞𝜙2 𝑧 ∣ 𝐷2
Variational Binary Selector: 𝜃1 𝑠, 𝑧
Variational Binary Selector: 𝜃2 𝑠, 𝑧
Preference Statistics: 𝜇1 , 𝜎1
Preference Statistics: 𝜇2 , 𝜎2
Model Aggregation
𝑝mixture 𝑧 = σ𝑖 𝑤𝑖 ⋅ 𝒩 𝑧; 𝜇𝑖 , 𝜎𝑖2 𝐼
𝜙𝑔 , 𝜃𝑔 ← FedAvg 𝜙𝑖 , 𝜃𝑖
𝑝2 (Harmless)
𝑧 Latent Space
Differentiable Prior
(c) Preference Alignment Reinforcement Learning
Global Server Federated Mixture Prior
𝑝1 (Helpful)
𝑧prior
𝑝mixture
Local Training
Orthogonal Loss: − cos 𝑧, 𝑝
Unlabeled Pairs 𝑥, 𝑠𝐴 , 𝑠𝐵
Labeling Process 𝑧 ~ 𝑝mixture 𝑦ො = argmax 𝜃𝑔 𝑠, 𝑧
Pseudo Labeled Dataset
Final Aligned LLM
Figure 1. Overview of the proposed FedVPA-GP framework. (a) Illustrates the federated training process of the variational binary selector. (b) Details the local variational objective designed to enhance inference. (c) Depicts the subsequent preference alignment stage using the trained selector.
the personalization mechanism ineffective in decentralized environments (Bowman et al., 2016; Alemi et al., 2018).
while preserving data privacy. • Stabilized Variational Inference and Disentanglement: To overcome data scarcity and prevent posterior collapse in federated settings, we introduce a mechanism combining a Federated Mixture Prior with an Orthogonal Loss. This approach stabilizes posterior estimation and enforces the semantic separation of distinct preference prototypes.
To overcome these challenges, we propose Federated Variational Preference Alignment with Gumbel-Softmax Prior (FedVPA-GP). We bridge the gap between local data sparsity and global distribution requirements through two core mechanisms. First, we introduce a Federated Mixture Prior that aggregates learned distributions from other clients, serving as a dynamic prior that stabilizes local inference. Second, to explicitly prevent posterior collapse and ensure semantic disentanglement, we incorporate an Orthogonal Loss that enforces the separation of conflicting preference prototypes in the latent space. By combining these with Gumbel-Softmax relaxation for end-to-end differentiability (Jang et al., 2017), FedVPA-GP successfully learns personalized reward models without sharing raw data, as shown in the Figure 2b.
• Empirical Validation: Experiments on the HH-RLHF dataset (Bai et al., 2022) demonstrate that FedVPA-GP significantly outperforms monolithic baselines (e.g., FedBiscuit, FedDPO) with robust generalization to unseen clients. Qualitative analysis confirms that our model successfully disentangles preferences.
2. Related Works
Extensive experiments on the HH-RLHF dataset (Bai et al., 2022) demonstrate that FedVPA-GP significantly outperforms monolithic baselines. Qualitative analysis further confirms that our algorithm successfully disentangles preferences in the latent space, enabling the model to dynamically switch between helpful and harmless modes based on the inferred context.
Reinforcement Learning from Human Feedback (RLHF) Since the seminal work of Christiano et al. (2017), RLHF has become a standard framework for aligning LLMs. The typical pipeline involves training a reward model on preference pairs to guide policy optimization via PPO (Schulman et al., 2017; Ouyang et al., 2022). Recently, methods such as Direct Preference Optimization (DPO) (Rafailov et al., 2023), IPO (Azar et al., 2024), and KTO (Ethayarajh et al., 2024) have been proposed to stabilize training by optimizing the policy directly without an explicit reward model. These approaches primarily operate in centralized settings, assuming access to aggregated datasets. Applying them to scenarios where data is distributed across edge devices introduces challenges related to data privacy and regula-
Our contributions are summarized as follows: • Federated Variational Preference Alignment: We address the limitation of monolithic reward models in capturing conflicting user preferences. By integrating variational inference into Federated Learning, our framework effectively adapts to diverse user intents 2
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
sA and sB are two candidate responses from a Large Language Model (LLM), and y ∈ {0, 1} indicates the user’s preference (with y = 1 denoting sA ≻ sB ). 3.1. Standard Federated Preference Alignment
(a) FedVPL
In standard Federated RLHF settings, the goal is to learn a global reward model rθ (sA , sB ) that maximizes the likelihood of user preferences across all clients. The preference probability is typically modeled using the Bradley-TerryLuce (BTL) model (Bradley & Terry, 1952; Luce, 1959):
(b) FedVPA-GP
pθ (y = 1|sA , sB ) = σ(rθ (sA ) − rθ (sB )),
Figure 2. Visualization of Latent Variable Distributions. (a) FedVPL suffers from posterior collapse, where latent codes zVPL cluster indistinguishably. (b) Our method (FedVPA-GP) effectively disentangles user preferences, showing distinct modes in zFedVPA-GP corresponding to different client groups.
(1)
where σ(·) is the sigmoid function. The federated objectivePminimizes the aggregate negative log-likelihood: K minθ i=1 EDi [− log pθ (y | sA , sB )]. However, this formulation assumes a single consensus reward function rθ , which inevitably averages out conflicting preferences (e.g., “Helpful” vs. “Harmless”) and fails to capture user-specific nuances (Poddar et al., 2024).
tory compliance (European Parliament and Council of the European Union, 2016). Federated Preference Alignment To address privacy concerns, recent studies have integrated alignment techniques with Federated Learning. Wu et al. (2024) proposed FedBiscuit, which utilizes client-side adapters to learn preference representations. Similarly, Ye et al. (2024) introduced FedDPO, extending DPO to the federated setting by aggregating gradients to update a global policy. These frameworks generally aim to learn a global consensus model. While effective for privacy, this global aggregation approach tends to average the preference distributions across clients, which may limit the model’s flexibility in scenarios where user preferences are heterogeneous or conflicting (Shirali et al., 2025).
3.2. Variational Preference Learning (VPL) To address heterogeneity, we adopt a latent conditional framework. We assume each user i is governed by a continuous latent preference vector zi ∈ Rd that conditions the reward model. For binary choice tasks, we condition the model’s logits on zi through a learned projection network: logits(sA , sB | zi ) = logitsbase (sA , sB ) + fθ (zi ),
(2)
where fθ : Rd → R|C| is a learned linear projection (latent projection) that maps the latent vector to logit adjustments, and C is the set of choices (typically {A, B}). The choice probability is then computed via softmax over the conditioned logits.
Personalized and Pluralistic Alignment Recognizing the diversity of human values (Santurkar et al., 2023), researchers have explored personalization in centralized settings. Techniques include multi-objective optimization (Rame et al., 2023), attribute steering (Dong et al., 2023), and weight merging (Jang et al., 2023). Notably, Variational Preference Learning (VPL) (Poddar et al., 2024) models user intent as a latent variable to capture continuous preference manifolds. However, these methods typically require access to the full dataset to learn the latent structure. Extending such variational approaches to Federated Learning presents specific challenges, particularly regarding local data sparsity and the estimation of stable posteriors in isolated environments.
Since zi is unobserved, we treat it as a latent variable and employ Variational Inference (VI) (Kingma & Welling, 2014). We introduce a local variational posterior qϕ (z | Di ) = N (z; µi , σi2 I) parameterized by ϕ to approximate the true posterior. The encoder extracts preference features (e.g., embedding difference ∆h = hchosen − hrejected ) and outputs posterior parameters (µi , σi2 ). The latent vector is sampled using the reparameterization trick: zi = µi + σi ⊙ ϵ, where ϵ ∼ N (0, I) and σi = exp(0.5·log σi2 ); ⊙ denotes elementwise multiplication. For brevity in the method (Sec. 4.1), we define qi := qϕ (z | Di ),
N0 := N (0, I).
(3)
Thus qi denotes client i’s variational posterior, and N0 the standard Gaussian prior.
3. Preliminaries
The objective is to maximize the Evidence Lower Bound (ELBO):
We consider a Federated Learning (FL) system consisting of K clients. Each client i ∈ {1, . . . , K} has a private dataset of pairwise preferences Di = {(sA , sB , y)}, where
L = Eqϕ (z|Di ) [log pθ (Di | z)] − βDKL (qi ∥ p(z)), (4) 3
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
where p(z) is the prior over latent preferences and β is a regularization coefficient. In the baseline VPL ablation we use p(z) = N0 .
the frozen base LLM. To isolate the preference signal from generic semantics, we construct a difference embedding ∆h = hchosen − hrejected from the response regions only. Concretely, we locate each response span via answer-token markers in the input and take the final-token hidden state of each span as its representation; this final-token pooling inherits attention from all preceding response tokens while excluding the prompt region.
3.3. Limitations of Federated Variational Preference Learning Transposing VPL to FL introduces critical challenges stemming from preference heterogeneity and data sparsity, which monolithic priors fail to address.
The difference vector ∆h is processed by a feature extractor (a multi-layer perceptron) that transforms the raw embedding difference into a lower-dimensional feature representation. This feature extractor learns to distill preferencespecific signals while suppressing general response characteristics. The processed features are then passed to the variational encoder to parameterize the local posterior distribution qϕ (z | Di ).
Sparsity and Instability: Data fragmentation leaves each client with a small local dataset |Di |, causing the variational posterior qi to be estimated with high variance. Insufficient samples lead to unstable gradients and poor convergence when training from scratch. To address this, we propose a Federated Mixture Prior, which leverages the aggregated distributions of other clients as a dynamic prior. This mechanism stabilizes local inference by transferring global knowledge, allowing clients to learn reliable posteriors even with sparse data.
qϕ (z | Di ) = N (z; µi , σi2 I).
(5)
To further guard against posterior collapse driven by unbounded variance, we cap the predicted log-variance, 2 log σi2 ← min(log σi2 , log σmax ); this prevents the encoder from trivially matching the prior by inflating σ. We then employ the reparameterization trick zi = µi + σi ⊙ ϵ, with ϵ ∼ N (0, I), to enable gradient-based optimization. By conditioning on ∆h, our design forces zi to encode the relative direction of user preferences rather than static response content.
Heterogeneity and Posterior Collapse: In centralized VPL, the model learns a global latent structure from pooled data. In FL, however, clients infer in isolation using a generic standard Gaussian prior N0 . This lack of global guidance often leads to posterior collapse, where the latent variable z degenerates to the uninformative prior and fails to encode personalized preferences (Bowman et al., 2016; Alemi et al., 2018). As illustrated in Figure 2b, this results in an entangled latent space where distinct preference clusters fail to emerge. To prevent this collapse and enforce a semantically meaningful structure, we introduce an Orthogonal Loss (Sec. 4.2), which explicitly separates conflicting preference prototypes.
Federated Mixture Prior with Learnable GumbelSoftmax Weights: To mitigate local data sparsity, we leverage the population-level distribution as a dynamic prior. However, simply averaging distributions from all clients is suboptimal due to preference heterogeneity. To address this, we propose a Federated Mixture Prior with learnable weights. Let S ⊆ {1, . . . , K} be the set of participat(i) ing clients. We construct the mixture prior pmixture (z) as a weighted sum of peer posteriors Nj (z): X (i) pmixture (z) = wj · Nj (z), (6)
4. Federated Variational Preference Alignment with Gumbel-Softmax Prior We propose Federated Variational Preference Alignment with Gumbel-Softmax Prior (FedVPA-GP), a framework designed to learn personalized reward models in a privacypreserving manner. Unlike previous approaches that simply aggregate gradients (Wu et al., 2024; Ye et al., 2024), FedVPA-GP treats user personalization as a distributed continuous latent variable inference problem. We first detail our variational inference mechanism with the proposed mixture prior, followed by the orthogonal regularization for disentanglement and finally the two-stage training strategy.
j∈S
where wj represents the relevance weight of client j’s distribution to the current client i. (i)
To compute the KL divergence DKL (qi ∥ pmixture ) stably, we employ the log-sum-exp trick for the log-mixture probability: X (i) log pmixture (z) = max aj +log exp(aj − max ak ) ,
4.1. Variational Inference with Federated Mixture Prior
j
j∈S
k
(7) where aj = log wj + log Nj (z). This formulation prevents numerical underflow when aggregating probabilities from numerous peers.
Inference Network: Each client maintains a local variational encoder qϕ (z | Di ). Given a preference pair (sA , sB , y), we extract hidden representations hA , hB from
4
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
local loss function for client i is:
Gumbel-Softmax Relaxation: Instead of static weighting, we optimize these coefficients to prioritize compatible peers using the Gumbel-Softmax relaxation (Jang et al., 2017). The weights are computed via the reparameterization trick:
X
Li (θ, ϕ) = −Ez∼qϕ
log pθ (y | sA , sB , z)
(sA ,sB ,y)∈Di
exp((log πj + gj )/τ ) wj = P , k∈S exp((log πk + gk )/τ )
|
{z
}
Lrecon
(8)
(i)
+ β · DKL (qϕ (z | Di ) ∥ pmixture (z)) | {z } Lreg (Prior Matching)
where πj are learnable logits, gj ∼ Gumbel(0, 1) is Gumbel noise, and τ is the temperature. By minimizing the KL divergence, the model automatically learns to upweight informative peers with similar preference structures while filtering out conflicting noise. The logits {πj } are local trainable parameters per client and are excluded from federated averaging, so each client retains a personalized peer-weighting strategy.
+ λ · Lorthogonal (z), | {z }
(10)
Lortho
where β controls KL regularization and λ weights the separation penalty. The first two terms constitute the negative ELBO, while the third enforces orthogonality. 4.4. Two-Stage Training Strategy
4.2. Orthogonal Loss for Preference Separation
Finally, we describe the deployment pipeline, adopting a two-stage strategy (Wu et al., 2024) to handle heterogeneity efficiently.
To ensure the latent space semantically separates diverse preference modes and prevents posterior collapse, we introduce an orthogonal loss motivated by (Li et al., 2024). We maintain a set of M learnable prototype vectors d {pm }M m=1 ⊂ R .
Stage 1 (Federated Selector Training): We train the variational binary preference selector using the objective Li (θ, ϕ) defined above. In this phase, each client learns a posterior qϕ (z | Di ) and predicts choices conditioned on zi . The preference prediction is performed via a Latent Conditional Reward Model:
Prototype Initialization: We initialize these prototypes using QR decomposition (Saxe et al., 2013). This process transforms a random initialization into a strictly orthonormal basis, ensuring that prototypes begin in mutually orthogonal subspaces. The resulting basis is then projected to a fixed radius to guarantee sufficient separation from the origin.
logits(sA , sB | zi ) = logitsbase (sA , sB ) + fθ (zi ),
Server-Side Label Assignment: To guide this separation, the server performs balanced k-means clustering (with k = M ) on the collected client means {µ̄i } from the previous round and assigns a prototype index yi∗ ∈ {1, . . . , M } to each client.
where fθ is a small MLP d→64→32→|C| mapping the latent vector zi to logit adjustments. Clients leverage the mixture prior for knowledge transfer, enabling stable inference despite local data sparsity without exchanging raw data.
Loss Computation: Clients encourage their latent z to align with the assigned prototype pyi∗ while maintaining orthogonality among all prototypes. The loss combines a pull term and an orthonormality constraint: Lorthogonal (z) = ∥z − pyi∗ ∥22 + γ · ∥PPT − IM ∥2F ,
(11)
Base-logit dropout. For base models where the frozen LLM already encodes a strong {A, B} preference signal (e.g., Qwen-2 0.5B), the latent residual fθ (z) above receives little gradient, exacerbating posterior collapse. We optionally apply Bernoulli dropout with rate plogit to the base choice-logit pair during training, forcing z to carry the full predictive signal on those steps. We use plogit = 0.5 for Qwen-2 0.5B and 0.0 for Gemma-2B.
(9)
where P ∈ RM ×d is the matrix of stacked prototypes. This mechanism forces latent representations into distinct orthogonal subspaces, effectively disentangling conflicting preferences (e.g., helpful vs. harmless).
Stage 2 (Conditional RLHF): We perform Centralized RLHF (Rafailov et al., 2023) on the server. We employ DPO to train a policy conditioned on the inferred client context z (e.g., z ∼ qi ). The converged selector from Stage 1 serves as the reward model, scoring generations as logits(sA , sB | z). This decoupling avoids the prohibitive communication costs of federated generation and mitigates training instability caused by conflicting local gradients (Wu et al., 2024).
4.3. Federated Variational Objective During Stage 1, we aim to maximize the Evidence Lower Bound (ELBO) regularized by the orthogonal loss. The 5
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
Algorithm 1 Server-Side: Federated Aggregation, Prior Management, and Stage 2 RLHF
Algorithm 2 Client-Side: Local Variational Training Require: Local dataset Di , global parameters θt , ϕt , mixture prior pmixture (z) (if t > 1), local steps E, learning rate η, KL weight β, orthogonal weight λ Ensure: Updated parameters θit , ϕti , average z distribution (µ̄i , σ̄i2 ), sample size ni 1: Receive θt , ϕt from server 2: if mixture prior received then 3: Update local prior: pmixture (z) ← {(µj , σj2 ), wj }j∈S t−1 4: else 5: Use standard prior: pmixture (z) = N (0, I) 6: end if 7: Initialize: θit ← θt , ϕti ← ϕt 8: Initialize: Zbatch ← ∅ (for collecting z values) 9: for local step e = 1, . . . , E do 10: for batch (sA , sB , y) ∈ Di do 11: Extract features: hchosen , hrejected ← LLM(sA , sB ) 12: Compute difference: ∆h = hchosen − hrejected 13: Process: fi ← FeatureExtractor(∆h) 14: Encode: (µi , σi2 ) ← qϕti (fi ) Sample: zi ∼ N (µi , σi2 I) via reparameterization trick 15: 16: Collect: Zbatch ← Zbatch ∪ {zi } 17: Project: ∆logits ← LatentProjection(zi ) 18: Condition: logits ← logitsbase + ∆logits Compute reconstruction loss: Lrecon ← − log pθit (y | 19: sA , sB , zi ) 20: Compute KL divergence: LKL ← β · DKL (qϕti (z | ·)∥pmixture (z)) 21: Compute orthogonal loss: Lortho ← λ · Lorthogonal (zi ) 22: Total loss: Li ← Lrecon + LKL + Lortho 23: Update: θit , ϕti via SGD on Li (only VPL components, base model frozen) 24: end for 25: end for 26: Compute average z distribution: µ̄i ← mean({µi }), σ̄i2 ← var({zi ∈ Zbatch }) 27: Send (θit , ϕti , µ̄i , σ̄i2 , |Di |) to server
Require: Clients K, rounds T , KL weight β, orthogonal weight λ Ensure: Global parameters θT , ϕT , client z distributions {µ̄Ti , σ̄i2 }K i=1 ; fine-tuned policy (Stage 2) 1: Stage 1: Federated selector training 2: Initialize θ0 , ϕ0 (base model frozen, only VPL components trainable) 3: for round t = 1, 2, . . . , T do 4: Sample clients S t ⊆ {1, . . . , K} (typically |S t | = 10) 5: Broadcast θt , ϕt to S t 6: if t > 1 then 7: Broadcast mixture prior {(µj , σj2 ), wj }j∈S t−1 to S t 8: end if 9: Wait for client updates t 10: Receive (θit , ϕti , µ̄i , σ̄i2 , nP i ) from each client i ∈ S t+1 t 1 11: Aggregate: θ ← |S t | i∈S t θi P 12: Aggregate: ϕt+1 ← |S1t | i∈S t ϕti Store mixture components {(µ̄i , σ̄i2 , ni )}i∈S t for next 13: round; the mixture weights wj are computed on the client side via Eq. 8 (learnable logits with Gumbel-Softmax relaxation). 14: end for 15: Stage 2: Conditional RL (selector as reward) 16: Load VPL components from selector: encoder qϕ , feature extractor, Z - TO - EMBEDDING 17: Load client average z {µ̄i }K i=1 (or compute from selector + training data) 18: Freeze selector (θT , ϕT ); use logits(sA , sB | z) as reward 19: for each DPO step (on server, no federated rounds) do 20: Get z for batch: from data, or µ̄i by client i, or infer via qϕ (· | features) 21: Inject z into policy: inputs embeds ← base embeds + Z - TO - EMBEDDING (z) 22: Generate conditioned on z; score with selector logits(sA , sB | z); update policy via DPO 23: end for
Stage 2 (Conditional RLHF) does not access any personal preference labels. Only prompts from the HH-RLHF corpus are used, and the (chosen, rejected) pairs consumed by DPO are constructed from on-policy generations of the current policy, labeled by the Stage-1 selector conditioned on the inferred client context z.
5. Experiments 5.1. Experimental Settings We evaluate our framework on the HH-RLHF dataset (Bai et al., 2022), which contains pairwise comparisons focused on helpfulness and harmlessness. To simulate a heterogeneous federated setting, we implement a strict Non-IID partition where clients are divided into two disjoint groups: 50% of clients exclusively hold preference pairs labeled for helpfulness, while the remaining 50% possess only harmlessness data. This partition models a scenario where local preference data are highly heterogeneous and conflicting. We vary the total number of clients K ∈ {10, 50, 100} and sample 5 clients per round for K = 10 and 10 clients per round for K ∈ {50, 100} to assess scalability. Accordingly, we set the number of orthogonal prototypes to M = 2 to match HH-RLHF’s two preference axes; M can be increased for richer preference spaces with more distinct user clusters, which we leave to future work.
Models We utilize two base models to validate performance across different scales: Qwen-2 0.5B (Yang et al., 2024) and Gemma-2B (Gemma Team et al., 2024). To ensure communication efficiency in the federated setting, both models are fine-tuned using LoRA (Hu et al., 2022). Detailed training configurations and hyperparameter settings are provided in the Appendix. Baselines We compare FedVPA-GP against three representative baselines: • FedDPO (Ye et al., 2024): Standard federated DPO. 6
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences Table 1. Main Results on HH-RLHF. We report the GPT-4 Win-rate (%) across varying client counts (N ∈ {10, 50, 100}). Shaded rows indicate our proposed method, FedVPA-GP, which consistently achieves the best trade-off between helpfulness and harmlessness.
10 C LIENTS
50 C LIENTS
100 C LIENTS
M ODEL
M ETHOD
H ELPFUL
H ARMLESS
H ELPFUL
H ARMLESS
H ELPFUL
H ARMLESS
Q WEN 2
F ED DPO F ED B ISCUIT F EDVPL F EDVPA-GP
48.12 ± 1.52 48.85 ± 1.41 62.24 ± 1.25 66.45 ± 1.12
77.34 ± 2.41 75.12 ± 2.28 84.56 ± 1.95 89.21 ± 1.68
43.05 ± 2.15 44.21 ± 1.98 54.18 ± 1.72 58.32 ± 1.45
69.22 ± 2.85 71.45 ± 2.61 78.12 ± 2.12 84.05 ± 1.94
41.48 ± 2.32 42.33 ± 2.11 53.05 ± 1.88 55.18 ± 1.55
67.15 ± 2.64 69.42 ± 2.45 77.34 ± 2.21 82.31 ± 2.05
G EMMA -2B
F ED DPO F ED B ISCUIT F EDVPL F EDVPA-GP
52.34 ± 1.75 51.65 ± 1.58 66.82 ± 1.34 73.21 ± 1.15
83.12 ± 2.55 82.45 ± 2.32 89.15 ± 2.05 96.34 ± 1.75
44.15 ± 2.31 46.21 ± 2.05 56.41 ± 1.84 64.48 ± 1.52
78.45 ± 2.92 78.12 ± 2.74 84.34 ± 2.31 95.12 ± 2.05
41.22 ± 2.58 43.44 ± 2.22 53.25 ± 1.95 60.15 ± 1.68
75.33 ± 3.12 76.05 ± 2.88 80.42 ± 2.45 92.45 ± 2.15
Table 2. Unseen Client Generalization Results. We report the GPT-4 Win-rate (%) for seen and unseen clients. Shaded rows indicate our proposed method.
Seen
Unseen
Method
Helpful Harmless
Helpful Harmless
FedDPO FedBiscuit FedVPL FedVPA-GP
46.35 47.32 56.23 65.28
47.27 47.62 49.25 63.16
78.62 79.25 83.82 94.25
methods often suffer from a severe trade-off: they tend to align the model towards harmlessness at the expense of helpfulness. This is particularly evident in the Qwen-2 experiments, where the helpfulness win-rate of baselines stagnates or even decreases as the focus shifts to harmlessness. In contrast, FedVPA-GP effectively disentangles these conflicting heterogeneous preferences by conditioning the reward model on client-specific latent variables. Consequently, our method achieves a Pareto improvement, securing significantly higher win-rates in both helpfulness and harmlessness compared to all baselines, demonstrating the efficacy of personalization in satisfying diverse user needs.
79.15 78.42 75.21 91.23
• FedBiscuit (Wu et al., 2024): A Federated preference alignment algorithm that trains light-weight binary selector through FL.
Robustness to Heterogeneity and Data Sparsity The experimental results also highlight the challenge of scaling in federated settings. As the number of clients increases, the amount of data each local client possesses becomes increasingly sparse and the aggregate distribution more heterogeneous. Table 1 demonstrates that the performance of baselines, and even the naive FedVPL, deteriorates notably under these conditions. While monolithic approaches struggle to maintain performance amidst this increased noise, FedVPA-GP exhibits robustness. By leveraging the Federated Mixture Prior to share distributional knowledge without sharing raw data, our approach maintains consistent and high alignment performance even in large-scale settings with high data sparsity, validating its stability in decentralized environments.
• FedVPL: A naive adaptation of VPL (Poddar et al., 2024) to FL using a fixed standard Gaussian prior N (0, I) without the orthogonal loss. Evaluation Metrics Following standard benchmarks (Bai et al., 2022), we employ GPT-4o (OpenAI, 2024) as a judge to evaluate the quality of responses generated by the finetuned models against a frozen baseline. We report the Win-rate (%) for both Helpfulness and Harmlessness on a held-out test set, assessing the model’s ability to satisfy conflicting user preferences. 5.2. Personalized Preference Alignment
5.3. Analysis
Table 1 presents the GPT-4 win-rates of FedVPA-GP compared to state-of-the-art federated baselines on the HHRLHF dataset across varying client scales.
Analysis of Latent Space Disentanglement To understand how the model represents conflicting preferences, Figure 3 visualizes the evolution of the latent preference distribution (z) of 5 clients preferring helpfulness and 5 clients preferring harmlessness using t-SNE (van der Maaten & Hinton, 2008). Red and blue points correspond to latent z inferred from clients prioritizing harmlessness and helpful-
Overcoming the Limits of Monolithic Reward Models As hypothesized, baselines relying on monolithic reward models (FedDPO and FedBiscuit) struggle to reconcile conflicting preference objectives. As shown in Table 1, these 7
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
Figure 3. Evolution of client-specific latent preference distributions (z) across training rounds for FedVPL (top row) and FedVPA-GP (bottom row). Points are colored by preference type: red (harmlessness) and blue (helpfulness). Star markers (∗) indicate orthogonal prototypes in FedVPA-GP. FedVPA-GP achieves better separation between preference types compared to FedVPL. 96.34
89.21 84.5685.7086.00
90
80 70
66.45 63.2464.1064.50
60
Win rate (%)
Win rate (%)
90
89.1590.2090.80
80
strategy where 5 clients from each cluster were used for training (Seen), while the remaining 5 clients from each cluster were reserved for evaluation (Unseen). For the variational approaches (FedVPL and FedVPA-GP), we performed variational inference on the unseen clients’ local datasets to estimate their latent preference vectors z without updating the model parameters. As shown in Table 2, monolithic baselines like FedDPO and FedBiscuit exhibit consistent performance across seen and unseen groups, but their overall win-rates are limited due to their inability to model personalization. In contrast, FedVPL suffers a significant performance degradation on unseen clients, indicating a failure to generalize the latent preference structure. However, FedVPA-GP demonstrates stability, maintaining high win-rates on unseen clients that are comparable to the seen clients. This result suggests that our proposed mixture prior and orthogonal regularization enable the model to learn a semantically meaningful and continuous latent space, allowing it to effectively capture and condition on the preferences of novel users via simple inference.
FedVPL FedVPL+Ortho FedVPL+GB Prior FedVPA-GP (ours) 75.21
70
66.8267.8068.20
60 Harmlessness
Helpfulness
(a) Qwen-2 0.5B
Harmlessness
Helpfulness
(b) Gemma-2B
Figure 4. Ablation study: helpfulness and harmlessness win rate (%). (a) Qwen-2 0.5B. (b) Gemma-2B. Methods: FedVPL, FedVPL+Ortho, FedVPL+GB Prior, FedVPA-GP.
ness, respectively. As observed in the top row, the baseline FedVPL suffers from posterior collapse, where the distributions for these distinct preference types remain entangled and non-separable throughout the training process. In contrast, FedVPA-GP demonstrates a clear trajectory towards disentanglement. Driven by the Federated Mixture Prior and Orthogonal Loss, our method progressively structures the latent space, resulting in a sharp separation between the two preference types. This structured latent topology confirms that the model successfully learns to distinguish between conflicting user intents, enabling dynamic adaptation to local contexts.
5.4. Ablation Study Component Contributions We analyze the impact of our key components in Figure 4. Adding the Orthogonal Loss (FedVPL+Ortho) consistently improves both metrics by preventing latent overlap, thereby mitigating posterior collapse. Meanwhile, the Federated Mixture Prior (FedVPL+GB Prior) stabilizes training against data sparsity by leveraging the aggregate population distribution as a dy-
Generalization to Unseen Clients To evaluate the robustness of our framework against new users, we conducted an experiment with 20 clients, equally divided into helpfulness and harmlessness clusters. We utilized a hold-out 8
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences Table 3. Robustness to Client Population Ratios. GPT-4 Win-rate (%) on HH-RLHF (Qwen-2 0.5B, N = 10 clients) under varying helpfulness/harmlessness client population splits. Shaded rows indicate our method, FedVPA-GP.
Ratio (H/Hm) Method
GP transmits the per-client mixture statistics (µ̄i , σ̄i2 ) ∈ R32 ×R32 . With FP32, this amounts to only 256 Bytes per client per round, which is negligible compared to a single LoRA adapter (on the order of MBs) or the gradient payload.
Helpful Harmless
70 / 30
FedBiscuit FedVPA-GP
49.12 68.12
72.13 87.14
30 / 70
FedBiscuit FedVPA-GP
45.34 65.56
75.52 89.14
80 / 20
FedBiscuit FedVPA-GP
51.24 68.25
70.24 87.23
20 / 80
FedBiscuit FedVPA-GP
44.15 64.88
78.21 90.32
Training Latency Each federated round of FedVPA-GP takes approximately 1.18× the wall-clock time of FedDPO under matched batch size and local-update steps. This modest overhead is incurred by the additional forward pass through the variational encoder and the KL and orthogonal loss terms, and is a worthwhile trade-off given the Pareto improvements demonstrated in Section 5. Taken together, the additional memory, compute, and communication costs introduced by FedVPA-GP are negligible relative to the scale of the base LLM, making the framework readily deployable in realistic federated settings without altering existing infrastructure budgets.
namic prior. Ultimately, the full FedVPA-GP framework achieves superior performance, demonstrating that the Mixture Prior ensures robust learning while the Orthogonal Loss enforces semantic disentanglement, yielding the best tradeoff between conflicting preferences.
6. Conclusion We present FedVPA-GP, a federated framework that learns personalized reward models without sharing raw preference data. Existing federated alignment methods enforce a monolithic reward that averages out conflicting user intents, while naive variational personalization in this setting suffers from posterior collapse driven by local data sparsity and heterogeneity. To address these challenges, we introduce a Federated Mixture Prior that leverages the aggregate population distribution as a dynamic prior, together with an Orthogonal Loss that explicitly structures the latent space.
Robustness to Client Population Ratios The 50/50 split in Table 1 is an idealized symmetric case; in practice, the relative frequency of preference modes across clients can vary substantially. We therefore stress-test FedVPA-GP under four asymmetric splits between helpfulness-preferring and harmlessness-preferring clients (70/30, 30/70, 80/20, and 20/80) on Qwen-2 0.5B with N = 10 clients. As reported in Table 3, FedVPA-GP consistently outperforms the FedBiscuit baseline by approximately 17–20 points in helpfulness and 12–17 points in harmlessness across all four ratios. The Federated Mixture Prior is distributionaware: even when one preference mode is heavily underrepresented, the learnable Gumbel-Softmax weights allow each client to upweight informative peers and avoid the minority mode being averaged out by the monolithic update.
Empirical results on HH-RLHF show that FedVPA-GP significantly outperforms monolithic baselines, disentangling conflicting user intents within a structured latent space and generalizing to unseen clients via inference alone. Our work provides a foundation for personalized, privacy-preserving LLM alignment. Future directions include extending this framework to capture more granular, multi-dimensional preference attributes and investigating its scalability in largescale cross-device settings.
5.5. Computational and Communication Efficiency We quantify the practical overhead introduced by FedVPAGP on Qwen-2 0.5B; all numbers are reported per client per communication round unless stated otherwise.
Acknowledgment
Parameter Overhead The variational modules (feature extractor, variational encoder, latent projection, and prototypes) add approximately 0.9M trainable parameters — only 0.18% of the 494M base-model parameters — which is comparable to the LoRA-adapter footprint already required by every federated baseline.
This work was supported by Institute of Information & Communications Technology Planning & Evaluation (IITP) grants funded by the Korea government (MSIT) (No. IITP-2026-RS-2024-00437866, Information Technology Research Center (ITRC); No. RS-2024-00509258, Global AI Frontier Lab; No. RS-2026-25511821, ITRC Development of Personalized Media Service Recommendation and Generative Technology; and No. RS-2019-II191906, Artificial Intelligence Graduate School Program (POSTECH)).
Communication Overhead In addition to the gradients and LoRA weights exchanged by all baselines, FedVPA9
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
Impact Statement
Ethayarajh, K., Xu, W., Muennighoff, N., Jurafsky, D., and Kiela, D. Kto: Model alignment as prospect theoretic optimization. In International Conference on Machine Learning, 2024. URL https://openreview.net/ forum?id=Duqy5E9nF8.
This work advances privacy-preserving AI by enabling the alignment of LLMs with diverse user preferences without centralizing sensitive data. By moving away from monolithic value systems, our framework respects the inherent pluralism of human values, allowing models to adapt to conflicting objectives like helpfulness and harmlessness. However, extreme personalization carries the risk of creating ”filter bubbles” where models might reinforce harmful user biases. While our method explicitly models harmlessness to mitigate this, future deployment must carefully balance personalization with robust safety guardrails to ensure ethical boundaries are maintained.
European Parliament and Council of the European Union. Regulation (eu) 2016/679 of the european parliament and of the council of 27 april 2016 on the protection of natural persons with regard to the processing of personal data and on the free movement of such data. Official Journal of the European Union, L119:1–88, 2016. Gemma Team, Mesnard, T., Hardin, C., Dadashi, R., Bhupatiraju, S., Pathak, S., Sifre, L., Rivière, M., Kale, M. S., Love, J., et al. Gemma: Open models. arXiv preprint arXiv:2403.08295, 2024.
References Alemi, A., Poole, B., Fischer, I., Dillon, J., Saurous, R. A., and Murphy, K. Fixing a broken ELBO. In International Conference on Machine Learning (ICML), pp. 159–168. PMLR, 2018.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. URL https:// openreview.net/forum?id=nZeVKeeFYf9.
Azar, M. G., Rowland, M., Piot, B., Guo, D., Calandriello, D., Valko, M., and Munos, R. A general theoretical paradigm to understand learning from human preferences. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2024.
Jang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations (ICLR), 2017.
Bai, Y., Jones, A., Ndousse, K., Askell, A., Chen, A., DasSarma, N., Drain, D., et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862, 2022.
Jang, J., Kim, S., Lin, B. Y., Wang, Y., Hessel, J., Zettlemoyer, L., Hajishirzi, H., Choi, Y., and Ammanabrolu, P. Personalized soups: Personalized large language model alignment via post-hoc parameter merging, 2023. URL https://arxiv.org/abs/2310.11564.
Bowman, S. R., Vilnis, L., Vinyals, O., Dai, A., Jozefowicz, R., and Bengio, S. Generating sentences from a continuous space. In Proceedings of the 20th SIGNLL Conference on Computational Natural Language Learning (CoNLL), pp. 10–21, 2016.
Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., et al. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning, 14(1–2):1–210, 2021. Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR), 2014.
Bradley, R. A. and Terry, M. E. Rank analysis of incomplete block designs: I. the method of paired comparisons. Biometrika, 39(3/4):324–345, 1952.
Li, H., Nguyen, M., and Pimentel-Alarcón, D. Preventing collapse in contrastive learning with orthonormal prototypes (clop), 2024. URL https://arxiv.org/ abs/2403.18699.
Carlini, N., Tramer, F., Wallace, E., Jagielski, M., HerbertVoss, A., Lee, K., Roberts, A., Brown, T., Song, D., Erlingsson, U., et al. Extracting training data from large language models. In USENIX Security Symposium, volume 6, 2021.
Luce, R. D. Individual choice behavior: A theoretical analysis. John Wiley & Sons, 1959.
Christiano, P. F., Leike, J., Brown, T., Milani, M., Amodei, D., and Amodei, D. Deep reinforcement learning from human preferences. In Advances in Neural Information Processing Systems (NeurIPS), volume 30, 2017.
McMahan, B., Moore, E., Ramage, D., Hampson, S., and Arcas, B. A. y. Communication-efficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics (AISTATS), pp. 1273–1282, 2017.
Dong, Y., Wang, Z., Sreedhar, M. N., Wu, X., and Kuchaiev, O. Steerlm: Attribute conditioned sft as an (usersteerable) alternative to rlhf, 2023. URL https:// arxiv.org/abs/2310.05344.
OpenAI. Gpt-4o system card, 2024. URL https:// openai.com/index/gpt-4o-system-card/. OpenAI. 10
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pp. 27730– 27744, 2022.
Liu, T., Ge, W., Deng, X., Zhou, X., Ren, X., Zhang, X., Wei, X., Ren, X., Liu, X., Fan, Y., Yao, Y., Zhang, Y., Wan, Y., Chu, Y., Liu, Y., Cui, Z., Zhang, Z., Guo, Z., and Fan, Z. Qwen2 technical report, 2024. URL https://arxiv.org/abs/2407.10671. Ye, R., Wang, W., Chai, J., Li, D., Li, Z., Xu, Y., Du, Y., Wang, Y., and Chen, S. Openfedllm: Training large language models on decentralized private data via federated learning. In Proceedings of the ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD), pp. 6137–6147, 2024.
Poddar, S., Wan, Y., Ivison, H., Gupta, A., and Jaques, N. Personalizing reinforcement learning from human feedback with variational preference learning. arXiv preprint arXiv:2408.10075, 2024. Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. In Advances in Neural Information Processing Systems, volume 36, 2023.
Ziegler, D. M., Stiennon, N., Wu, J., Brown, T. B., Radford, A., Amodei, D., and Christiano, P. F. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.08593, 2019.
Rame, A., Couairon, G., Dancette, C., Gaya, J.-B., Shukor, M., Soulier, L., and Cord, M. Rewarded soups: towards pareto-optimal alignment by interpolating weights finetuned on diverse rewards. Advances in Neural Information Processing Systems, 36:71095–71134, 2023. Santurkar, S., Durmus, E., Ladhak, F., Lee, C., Liang, P., and Hashimoto, T. Whose opinions do language models reflect? In International Conference on Machine Learning (ICML), 2023. Saxe, A. M., McClelland, J. L., and Ganguli, S. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. Shirali, A., Nasr-Esfahany, A., Alomar, A., Mirtaheri, P., Abebe, R., and Procaccia, A. D. Direct alignment with heterogeneous preferences. arXiv preprint arXiv:2502.16320, 2025. van der Maaten, L. and Hinton, G. Visualizing data using t-sne. Journal of Machine Learning Research, 9(11): 2579–2605, 2008. Wu, F., Liu, X., Wang, H., Wang, X., and Gao, J. Towards federated rlhf with aggregated client preference for llms. arXiv preprint arXiv:2407.03038, 2024. Yang, A., Yang, B., Hui, B., Zheng, B., Yu, B., Zhou, C., Li, C., Li, C., Liu, D., Huang, F., Dong, G., Wei, H., Lin, H., Tang, J., Wang, J., Yang, J., Tu, J., Zhang, J., Ma, J., Yang, J., Xu, J., Zhou, J., Bai, J., He, J., Lin, J., Dang, K., Lu, K., Chen, K., Yang, K., Li, M., Xue, M., Ni, N., Zhang, P., Wang, P., Peng, R., Men, R., Gao, R., Lin, R., Wang, S., Bai, S., Tan, S., Zhu, T., Li, T., 11
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
A. Mathematical Proofs This section provides detailed mathematical proofs and derivations for key components of our method. A.1. KL Divergence with Mixture Prior A.1.1. D EFINITION AND C OMPUTATION Theorem 1 (KL Divergence with Mixture Prior): The KL divergence between a posterior distribution qi (z) = N (z; µi , σi2 I) P|S| and a mixture prior pmixture (z) = j=1 wj · N (z; µj , σj2 I) is given by: DKL (qi ∥ pmixture ) = Ez∼qi [log qi (z) − log pmixture (z)] ,
(12)
where S is the set of participating clients from the previous round and wj are the Gumbel-Softmax mixture weights computed on the client side via Eq. 8 from the local learnable logits πj . Proof: By definition of KL divergence:
Z DKL (qi ∥ pmixture ) =
qi (z) log
qi (z) dz pmixture (z)
Z =
qi (z) [log qi (z) − log pmixture (z)] dz
= Ez∼qi [log qi (z) − log pmixture (z)] .
(13)
For a multivariate Gaussian posterior qi (z) = N (z; µi , σi2 I) with dimension d, we have: d d d 1X 1 X (zj − µi,j )2 2 log qi (z) = − log(2π) − log σi,j − . 2 2 2 j=1 2 j=1 σi,j
(14)
For the mixture prior, we compute: |S| X log pmixture (z) = log wj · N (z; µj , σj2 I) .
(15)
j=1
Using Monte Carlo estimation with batch size B: B i 1 Xh DKL ≈ log qi (z (b) ) − log pmixture (z (b) ) , B
(16)
b=1
where z (b) ∼ qi is sampled using the reparameterization trick: z (b) = µi + σi ⊙ ϵ(b) ,
ϵ(b) ∼ N (0, I),
(17)
and σi = exp(0.5 · log σi2 ) with ⊙ denoting element-wise multiplication. □ A.1.2. L OG -S UM -E XP T RICK FOR N UMERICAL S TABILITY Theorem 2 (Log-Sum-Exp Trick): For numerical stability when computing log pmixture (z), we use the log-sum-exp trick: 12
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
log
N X
! exp(ai )
= max ai + log i
i=1
N X
! exp(ai − max ai ) . i
i=1
(18)
Proof: We factor out the maximum term: N X
exp(ai ) = exp(max ai ) · i
i=1
N X
exp(ai − max ai ).
(19)
i
i=1
Taking the logarithm of both sides:
log
N X
! exp(ai )
= max ai + log
i=1
i
N X
! exp(ai − max ai ) . i
i=1
(20)
Since exp(ai − maxi ai ) ∈ [0, 1] for all i, this formulation is numerically stable and prevents overflow/underflow. For the mixture prior, we apply this trick with: aj = log wj + log N (z; µj , σj2 I),
(21)
where: d
d
k=1
k=1
1 X (zk − µj,k )2 1X d 2 log σj,k − . log N (z; µj , σj2 I) = − log(2π) − 2 2 2 2 σj,k
(22)
□ A.2. Reparameterization Trick and Gradient Flow Theorem 3 (Reparameterization Trick Gradient Flow): Using the reparameterization trick, gradients with respect to z flow to µ and log σ 2 . Proof: The reparameterization trick expresses the random variable z as a deterministic function of parameters and noise: z = µ + ϵ ⊙ exp(0.5 · log σ 2 ),
ϵ ∼ N (0, I),
(23)
where σ = exp(0.5 · log σ 2 ). The partial derivatives are: ∂z = I (identity matrix), ∂µ ∂z = ϵ ⊙ exp(0.5 · log σ 2 ) ⊙ 0.5. ∂ log σ 2
(24) (25)
By the chain rule, for any function f (z): ∂f (z) ∂f (z) ∂z ∂f (z) = · = , ∂µ ∂z ∂µ ∂z ∂f (z) ∂f (z) ∂z ∂f (z) = · = · ϵ ⊙ σ ⊙ 0.5. 2 2 ∂ log σ ∂z ∂ log σ ∂z 13
(26) (27)
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
This ensures that gradients can flow through the sampling operation, enabling end-to-end training of the variational encoder. □ A.3. Orthogonal Loss Formulation A.3.1. P ULL L OSS The pull loss encourages latent representations z to align with their assigned prototypes: Lpull (z) = ∥z − pyi∗ ∥22 ,
(28)
where pyi∗ is the prototype assigned to client i based on the server’s orthogonal label yi∗ ∈ {0, 1}. A.3.2. O RTHONORMALITY C ONSTRAINT To maintain orthogonality between prototypes, we enforce an orthonormality constraint: Lorthonorm = ∥PT P − I∥2F ,
(29)
where P = [p0 , p1 ]T is the prototype matrix and ∥ · ∥F denotes the Frobenius norm. This constraint ensures that pT0 p1 = 0 (orthogonality) and ∥p0 ∥2 = ∥p1 ∥2 = 1 (normalization). A.3.3. T OTAL O RTHOGONAL L OSS The combined orthogonal loss is: Lorthogonal (z) = λ · Lpull (z) + γ · Lorthonorm ,
(30)
where λ = 1.0 (orthogonal weight) and γ = 0.1 (orthonorm weight) are hyperparameters. This loss encourages latent representations to cluster around their assigned prototypes while ensuring that different preference types occupy orthogonal subspaces, thereby suppressing general features that are not preference-specific and preventing neural collapse. A.4. Evidence Lower Bound (ELBO) Derivation Theorem 4 (ELBO Derivation): The Evidence Lower Bound (ELBO) for variational inference is: LELBO = Eqϕ (z|Di ) [log pθ (Di | z)] − β DKL (qϕ (z | Di ) ∥ p(z)),
(31)
where β is a regularization coefficient. Proof: We start with the log-likelihood of the data: Z pθ (Di | z)p(z)dz.
log pθ (Di ) = log
(32)
Introducing the variational posterior qϕ (z | Di ): pθ (Di | z)p(z) dz qϕ (z | Di ) · qϕ (z | Di ) pθ (Di | z)p(z) = log Eqϕ (z|Di ) . qϕ (z | Di ) Z
log pθ (Di ) = log
14
(33)
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
Applying Jensen’s inequality (since log is concave): pθ (Di | z)p(z) log pθ (Di ) ≥ Eqϕ (z|Di ) log qϕ (z | Di ) = Eqϕ (z|Di ) [log pθ (Di | z)] + Eqϕ (z|Di ) log
p(z) qϕ (z | Di )
= Eqϕ (z|Di ) [log pθ (Di | z)] − DKL (qϕ (z | Di ) ∥ p(z)).
(34)
Following the β-VAE formulation (Alemi et al., 2018), we re-weight the KL term by a regularization coefficient β to control prior-matching pressure: LELBO = Eqϕ (z|Di ) [log pθ (Di | z)] − β DKL (qϕ (z | Di ) ∥ p(z)).
(35)
For β ̸= 1, the resulting objective is no longer a strict lower bound on log pθ (Di ); it instead trades reconstruction fidelity against KL pressure, a regime known to mitigate posterior collapse in low-data settings (Alemi et al., 2018; Bowman et al., 2016). The first term is the reconstruction loss (preference alignment), and the second term is the regularization (prior matching). Maximizing the ELBO is equivalent to minimizing the negative ELBO: Li (θ, ϕ) = −ELBO(Di ) = −Eqϕ [log pθ (Di | z)] + β DKL (qϕ (z | Di ) ∥ p(z)).
(36)
□ A.5. Standard Gaussian Prior KL Divergence For the baseline VPL ablation, we use a standard Gaussian prior p(z) = N (0, I). The KL divergence has a closed-form expression: Theorem 5 (Standard Gaussian Prior KL): For qi (z) = N (z; µi , σi2 I) and p(z) = N (0, I), the KL divergence is:
DKL (qi ∥ N0 ) =
d 1X 2 2 2 µi,j + σi,j − log σi,j −1 . 2 j=1
(37)
Proof: For two multivariate Gaussians, the KL divergence is: 1 |Σ2 | −1 T −1 DKL (N (µ1 , Σ1 ) ∥ N (µ2 , Σ2 )) = tr(Σ2 Σ1 ) + (µ2 − µ1 ) Σ2 (µ2 − µ1 ) − d + log . 2 |Σ1 |
(38)
For qi = N (µi , σi2 I) and p = N (0, I): |I| 1 −1 2 T −1 DKL (qi ∥ N0 ) = tr(I · σi I) + (0 − µi ) I (0 − µi ) − d + log 2 2 |σi I| 1 2 T 2 = tr(σ I) + µi µi − d − log |σi I| 2 i =
d d d X X 1 X 2 2 σi,j + µ2i,j − d − log σi,j 2 j=1 j=1 j=1
=
d 1X 2 2 2 µ + σi,j − log σi,j −1 . 2 j=1 i,j
15
(39)
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
□ A.6. Gumbel-Softmax for Differentiable Prior Sampling For sampling from the mixture prior (used in visualization and generation), we employ Gumbel-Softmax relaxation (Jang et al., 2017) with temperature τ = 1.0 to enable differentiable sampling. Component probabilities: exp((log wj + gj )/τ ) αj = P|S| , k=1 exp((log wk + gk )/τ )
(40)
where gj ∼ Gumbel(0, 1) are independent Gumbel random variables. Sampling:
zprior =
|S| X
where zj ∼ N (µj , σj2 I).
αj · zj ,
(41)
j=1
As τ → 0, this approaches categorical sampling (hard assignment), while τ > 0 provides a smooth, differentiable approximation. A.7. Stop-Gradient on Peer-Provided Prior Parameters Remark (Stop-Gradient on Peer Statistics): Within a single client’s local update, the mixture-prior parameters {µj , σj2 }j∈S provided by peer clients enter the computation graph as detached constants and therefore receive no gradient on that client’s pass. Justification: For the current client i, the mixture prior (i)
pmixture (z) =
X
wj · N (z; µj , σj2 I)
(42)
j∈S
is constructed from peer means µj and variances σj2 that were computed during round t − 1 and broadcast to client i. These tensors are not leaf nodes in client i’s autograd graph, so ∂Li ∂Li = =0 ∂µj ∂σj2
(43)
on client i’s backward pass. Peer parameters are updated by their respective owning clients in their own local training rounds; the Gumbel-Softmax weights wj , in contrast, are computed from the client-local trainable logits πj (Eq. 8) and therefore do receive gradients. (i)
In practice, the gradient of Lrecon +βDKL (qi ∥ pmixture )+Lortho with respect to the local variational parameters (µi , log σi2 ) is obtained by automatic differentiation through the reparameterized sample zi = µi + σi ⊙ ϵ and the log-sum-exp computation (i) of log pmixture (zi ); no manual derivation is required. □
B. Hyperparameter Details B.1. Hyperparameter Settings Table 4 provides the final hyperparameter values used in our experiments. 16
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences Table 4. Final hyperparameter settings for selector training (Stage 1) and RL training (Stage 2). Parameter Learning rate Batch size Gradient accumulation steps Local update steps Total rounds KL weight (β) Orthogonal loss weight (λ) Orthonorm weight (γ) Gumbel-Softmax temperature (τ ) Prototype scale (s) Latent dimension (d) Latent projection fθ Orthogonal label assignment LoRA rank (r) LoRA alpha (α) LoRA dropout (p) Reward coefficient Max prompts for generation Generation batch size Max samples for reward
Selector Training −5
−4
10 to 10 (model-dependent) 4–8 (model-dependent) 4–8 30 50 0.01 1.0 0.1 1.0 5.0 32 MLP d→64→32→|C| Balanced k-means 8 16 0.05 – – – –
RL Training 10−5 to 10−4 (model-dependent) 1 4–32 (model-dependent) 30 50 – – – – – – – – 8 16 0.05 0.1 50 3 30
C. Experimental Settings C.1. Dataset Details C.1.1. HH-RLHF DATASET We use the HH-RLHF (Helpful and Harmless from Human Feedback) dataset (Bai et al., 2022), which contains pairwise preference comparisons along two axes: helpfulness and harmlessness. Data splits: • Train: 80% • Validation: 10% • Test: 10% Client configurations: • Number of clients: K ∈ {10, 50, 100} • Sampling rates per round: – K = 10: 5 clients per round – K ∈ {50, 100}: 10 clients per round Data characteristics: • Data type: Pairwise preference comparisons • Preference axes: Helpfulness, Harmlessness • Each sample: (sA , sB , y) where y ∈ {0, 1} indicates preference 17
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
C.2. Model Details C.2.1. BASE L ANGUAGE M ODELS We conduct experiments using two base language models: • Qwen-2 0.5B: A compact 0.5 billion parameter model from the Qwen-2 family (?), suitable for resource-constrained federated environments. • Gemma-2B: A 2 billion parameter model from Google’s Gemma family (Gemma Team et al., 2024), providing a larger model baseline for comparison. C.2.2. F INE - TUNING C ONFIGURATION Both models are fine-tuned using LoRA (Low-Rank Adaptation) (Hu et al., 2022) to enable parameter-efficient fine-tuning in federated settings. LoRA parameters: • LoRA rank: r = 8 • LoRA alpha: α = 16 • Dropout rate: p = 0.05 Training configuration: • During federated selector training (Stage 1): The base LLM is frozen; only the VPL components (feature extractor, variational encoder, latent projection, and orthogonal prototypes) and LoRA adapters are updated. • During RL training (Stage 2): The base LLM remains frozen; LoRA adapters and the Z - TO - EMBEDDING module are trained via DPO conditioned on the inferred client context z, with the Stage 1 selector providing reward signals. C.3. Evaluation Settings C.3.1. W INRATE E VALUATION We evaluate the final policy using GPT-4 Win-rate (Bai et al., 2022): • Judge: GPT-4 or gpt-4o-mini (for cost efficiency) • Comparison: Fine-tuned model response vs baseline (frozen) model response • Evaluation samples: Up to 30 samples per evaluation (for efficiency) • Metrics: – Helpful Win-rate (%): Percentage of responses where fine-tuned model is more helpful – Harmless Win-rate (%): Percentage of responses where fine-tuned model is more harmless C.3.2. R EWARD M ODEL E VALUATION We also evaluate using reward model scores: • Harmlessness score: Reward model’s harmlessness score (higher is better) • Helpfulness score: Reward model’s helpfulness score (higher is better) • Evaluation samples: Up to 100 samples 18
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
C.4. Baseline Methods We compare our proposed FedVPA-GP with the following baselines: • FedDPO: Federated Direct Preference Optimization (Ye et al., 2024), which aggregates gradients for a single global policy without personalized reward modeling. • FedBiscuit: Federated learning with multiple LoRA adapters (U = 3) for coarse-grained personalization (Wu et al., 2024). • FedVPL: Our naive adaptation of VPL (Poddar et al., 2024) to FL—same latent selector with FedAvg, but using fixed Gaussian prior N (0, I) and no orthogonal loss. • FedVPA-GP (ours): Our full method with federated mixture prior, Gumbel-Softmax relaxation, difference embeddings, and orthogonal loss for preference separation.
D. Generation Prompts This section details all prompts used throughout our experiments for data generation, selection, and evaluation. D.1. HH-RLHF Comparison Prompt D.1.1. S ELECTOR T RAINING (B INARY C HOICE ) For binary selector training, we use the following prompt template to format preference comparison data: Below is a conversation between a human and an AI assistant, followed by two responses. Pick the response that is more helpful and harmless. State your choice with a single capital letter, i.e., "A" if RESPONSE A is better, "B" if RESPONSE B is better. ### CONVERSATION: {prompt} ### RESPONSE A: {output_A} ### RESPONSE B: {output_B} ### YOUR CHOICE: This prompt is used to train the binary preference selector, where the model learns to predict which response (A or B) is preferred based on helpfulness and harmlessness criteria. D.1.2. RL G ENERATION (S TANDARD G ENERATION ) For RL training, we generate responses using the conversation history as the prompt. The generation process uses the following settings: Generation parameters: • top p: 1.0 • temperature: 0.7 • do sample: True • max new tokens: 512 (configurable) • num return sequences: 2 (default) Prompt format: The prompt consists of the conversation history (all dialogue turns before the final assistant response). The model generates continuations from this prompt. 19
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
D.2. GPT API Winrate Evaluation Prompt For winrate evaluation using GPT API, we use the following prompt template to compare two responses: Below is a conversation between a human and an AI assistant, followed by two responses. Pick the response that is more helpful and harmless. State your choice with a single capital letter, i.e., "A" if RESPONSE A is better, "B" if RESPONSE B is better. ### CONVERSATION: {prompt} ### RESPONSE A: {response_a} ### RESPONSE B: {response_b} ### YOUR CHOICE: Evaluation process: 1. Generate responses from fine-tuned model for test prompts 2. Generate responses from baseline model (adapter disabled) for the same prompts 3. For each prompt, send the comparison prompt to GPT API (gpt-4o-mini by default) 4. Parse GPT response to extract choice (A or B) 5. Calculate winrate: percentage of cases where fine-tuned model (RESPONSE A) is preferred Configuration: • use gpt api for winrate: True • openai model: "gpt-4o-mini" (default, cost-efficient) • max samples for reward: 30 (evaluation sample limit) D.3. Additional Generation Prompts D.3.1. H ELPFULNESS -F OCUSED G ENERATION For helpfulness-specific generation (used in ablation studies): Below is a conversation between a human and an AI assistant. Write a response that is helpful. ### CONVERSATION: {prompt} ### RESPONSE: D.3.2. H ARMLESSNESS -F OCUSED G ENERATION For harmlessness-specific generation (used in ablation studies): Below is a conversation between a human and an AI assistant. Write a response that is harmless.
20
Federated Variational Preference Alignment with Gumbel-Softmax Prior for Personalized User Preferences
### CONVERSATION: {prompt} ### RESPONSE: D.3.3. G ENERAL G ENERATION For general response generation (both helpful and harmless): Below is a conversation between a human and an AI assistant. Write a response that is both helpful and harmless. ### CONVERSATION: {prompt} ### RESPONSE: D.4. Prompt Usage Summary Table 5 summarizes when each prompt template is used. Table 5. Prompt template usage across different stages of training and evaluation. Stage
Prompt Template
Selector Training RL Generation GPT Winrate Evaluation Helpfulness Ablation Harmlessness Ablation
Comparison prompt (binary choice) Conversation history (standard generation) Comparison prompt (A vs B) Helpfulness-focused generation Harmlessness-focused generation
21