HyperAIHyperAI

Command Palette

Search for a command to run...

MiCRo: Mixture Modeling and Context-aware Routing for Personalized Preference Learning

Jingyan Shen Jiarui Yao Rui Yang Yifan Sun Feng Luo Rui Pan Tong Zhang Han Zhao

Abstract

Reward modeling is a key step in building safe foundation models when applying reinforcement learning from human feedback (RLHF) to align Large Language Models (LLMs). However, reward modeling based on the Bradley Terry (BT) model assumes a global reward function, failing to capture the inherently diverse and heterogeneous human preferences. Hence, such oversimplification limits LLMs from supporting personalization and pluralistic alignment. Theoretically, we show that when human preferences follow a mixture distribution of diverse subgroups, a single BT model has an irreducible error. While existing solutions, such as multi-objective learning with finegrained annotations, help address this issue, they are costly and constrained by predefined attributes, failing to fully capture the richness of human values. In this work, we introduce MiCRo, a two-stage framework that enhances personalized preference learning by leveraging large-scale binary preference datasets without requiring explicit fine-grained annotations. In the first stage, MiCRo introduces context-aware mixture modeling approach to capture diverse human preferences. In the second stage, Mi CRo integrates an online routing strategy that dynamically adapts mixture weights based on specific context to resolve ambiguity, allowing for efficient and scalable preference adaptation with minimal additional supervision. Experiments on multiple preference datasets demonstrate that MiCRo effectively captures diverse human preferences and significantly improves downstream personalization.

One-sentence Summary

Researchers from University of Illinois Urbana-Champaign, New York University, and Rice University propose MiCRo, a two-stage framework that combines context-aware mixture modeling with an online routing strategy to capture diverse human preferences from large-scale binary feedback without requiring fine-grained annotations, addressing the irreducible error inherent to single Bradley-Terry models and enabling personalized, pluralistic alignment in LLMs.

Key Contributions

  • A theoretical lower bound shows that a single Bradley-Terry reward model has irreducible cross-entropy error when human preferences follow a mixture distribution over diverse subgroups, with the error increasing as subgroup diversity and context-dependent subgroup weights increase.

  • MiCRo is a two-stage framework that leverages large-scale binary preference datasets without requiring explicit fine-grained annotations by first applying context-aware mixture modeling over multiple Bradley-Terry reward components and then using an online routing strategy to adapt mixture weights from prompt context.

  • Experiments on multiple preference datasets show that MiCRo better captures heterogeneous human preferences and improves downstream personalization and pluralistic alignment compared with approaches based on a single global reward model.

Introduction

Reinforcement learning from human feedback is used to align language models with human values, especially where preferences are subjective and open ended. Existing preference learning methods often assume a single global reward function from binary pairwise labels, which cannot represent diverse or conflicting user preferences. Recent attempts to model multiple attributes either require costly fine-grained annotations, rely on imperfect LLM-as-a-judge labels, or need manually designed personalization contexts. The authors introduce MiCRo, a two-stage context-aware mixture modeling framework that learns latent subpopulation reward functions from large-scale binary preference data, provides a theoretical lower bound showing a single Bradley-Terry model has irreducible error under mixture preferences, and uses context-aware online routing to adapt to personalized users from limited samples.

Method

The authors leverage a two-stage preference modeling framework that separates offline discovery of preference heterogeneity from online adaptation to user context. The first stage learns a set of candidate reward functions and a context-aware mixture weight network from a large-scale binary preference dataset. The second stage keeps the learned reward heads fixed or lightly updated and fine-tunes only a lightweight routing network to assign user-specific weights over the reward heads. This design allows the system to avoid costly annotation of subpopulation labels and to adapt to personalized preferences at deployment time.

At the core of the first stage is a mixture of Bradley-Terry reward models. In the standard single-reward setting, for a prompt xxx and two responses awa_waw​ and ala_lal​, the probability that awa_waw​ is preferred over ala_lal​ is modeled as:

P(aw≻al∣x)=σ(r∗(x,aw)−r∗(x,al)),\mathbb{P}(a_w \succ a_l \mid x) = \sigma(r^*(x,a_w)-r^*(x,a_l)),P(aw​≻al​∣x)=σ(r∗(x,aw​)−r∗(x,al​)),

where σ(⋅)\sigma(\cdot)σ(⋅) is the logistic function and r∗r^*r∗ is a latent reward function. To capture diverse human preferences without additional annotation, the authors introduce KKK latent subpopulations, each associated with its own reward function rϕkr_{\phi_k}rϕk​​. They also introduce a mixture weight network fψf_\psifψ​ that maps each prompt xxx to a probability vector over the KKK experts:

fψ(x)∈ΔK,∑k=1Kfψ,k(x)=1,fψ,k(x)≥0.f_\psi(x) \in \Delta^K, \qquad \sum_{k=1}^K f_{\psi,k}(x)=1, \qquad f_{\psi,k}(x)\geq 0.fψ​(x)∈ΔK,k=1∑K​fψ,k​(x)=1,fψ,k​(x)≥0.

Using this mixture weight network, the overall preference probability becomes:

P(aw≻al∣x)=∑k=1Kfψ,k(x) σ(rϕk(x,aw)−rϕk(x,al)).\mathbb{P}(a_w \succ a_l \mid x) = \sum_{k=1}^K f_{\psi,k}(x) \, \sigma\left(r_{\phi_k}(x,a_w)-r_{\phi_k}(x,a_l)\right).P(aw​≻al​∣x)=k=1∑K​fψ,k​(x)σ(rϕk​​(x,aw​)−rϕk​​(x,al​)).

This formulation makes the mixture context-aware. Instead of averaging all reward heads with fixed weights, the model allows each prompt to activate the most relevant expert or experts. The architecture therefore contains KKK reward heads and one prompt-conditioned router, where each reward head represents a latent preference profile and the router decides how strongly each profile should contribute to the final prediction for a given prompt.

The reward heads and router are trained jointly by minimizing a negative log-likelihood loss over the preference dataset:

Lmle=−1n∑(x,aw,al)∈Dlog⁡∑k=1K(fψ,k(x) σ(rϕk(x,aw)−rϕk(x,al))).\mathcal{L}_{\mathrm{mle}} = -\frac{1}{n} \sum_{(x,a_w,a_l)\in\mathcal{D}} \log \sum_{k=1}^K \Big( f_{\psi,k}(x) \, \sigma\left(r_{\phi_k}(x,a_w)-r_{\phi_k}(x,a_l)\right) \Big).Lmle​=−n1​(x,aw​,al​)∈D∑​logk=1∑K​(fψ,k​(x)σ(rϕk​​(x,aw​)−rϕk​​(x,al​))).

To prevent a single reward head from dominating the mixture, the authors add a regularization term that encourages the router to keep its weight distribution non-degenerate:

Lreg=1n∑(x,aw,al)∈D∑k=1Kfψ,k(x)log⁡fψ,k(x).\mathcal{L}_{\mathrm{reg}} = \frac{1}{n} \sum_{(x,a_w,a_l)\in\mathcal{D}} \sum_{k=1}^K f_{\psi,k}(x)\log f_{\psi,k}(x).Lreg​=n1​(x,aw​,al​)∈D∑​k=1∑K​fψ,k​(x)logfψ,k​(x).

The full first-stage objective is:

L(ϕ,ψ)=Lmle+αLreg,\mathcal{L}(\phi,\psi) = \mathcal{L}_{\mathrm{mle}} + \alpha \mathcal{L}_{\mathrm{reg}},L(ϕ,ψ)=Lmle​+αLreg​,

where α\alphaα controls the strength of the regularization. This stage learns a diverse set of preference-specialized reward functions from standard binary preference labels. It does not require annotator identities, user profiles, or explicit subpopulation labels during training.

The second stage addresses the remaining deployment challenge: given a specific user or interaction context, how should the system assign weights to the previously learned reward heads? The authors handle this by fine-tuning the router using contextual data. In this stage, the router input expands from a prompt xix_ixi​ to a prompt-context pair (xi,ci)(x_i,c_i)(xi​,ci​), where cic_ici​ can include instructions, metadata, demographics, interaction history, or other user signals. The router now produces weights as:

fψ(xi,ci)∈ΔK.f_\psi(x_i,c_i)\in\Delta^K.fψ​(xi​,ci​)∈ΔK.

Each reward head remains available as an expert that evaluates the preference pair independently. For a contextual preference example (xi,ci,ai,w,ai,l)(x_i,c_i,a_{i,w},a_{i,l})(xi​,ci​,ai,w​,ai,l​), the probability assigned by the kkk-th expert is:

pk(ai,w≻ai,l∣xi)=σ(rϕk(xi,ai,w)−rϕk(xi,ai,l)).p_k(a_{i,w}\succ a_{i,l}\mid x_i) = \sigma\left(r_{\phi_k}(x_i,a_{i,w})-r_{\phi_k}(x_i,a_{i,l})\right).pk​(ai,w​≻ai,l​∣xi​)=σ(rϕk​​(xi​,ai,w​)−rϕk​​(xi​,ai,l​)).

A head-level preference loss is defined as:

Li,k=−log⁡pk(ai,w≻ai,l∣xi).\mathcal{L}_{i,k} = -\log p_k(a_{i,w}\succ a_{i,l}\mid x_i).Li,k​=−logpk​(ai,w​≻ai,l​∣xi​).

The router assigns weights to experts, and the second-stage objective combines the expected expert loss with a KL divergence proximal term:

Li=∑k=1Kfψ,k(xi,ci)Li,k+τ KL(fψ(xi,ci) ∥ ωi),\mathcal{L}_i = \sum_{k=1}^K f_{\psi,k}(x_i,c_i)\mathcal{L}_{i,k} + \tau \, \mathrm{KL} \left( f_\psi(x_i,c_i) \,\middle\|\, \omega_i \right),Li​=k=1∑K​fψ,k​(xi​,ci​)Li,k​+τKL(fψ​(xi​,ci​)∥ωi​),

where ωi\omega_iωi​ is a prior or previous weight vector and τ\tauτ controls how strongly the updated router should stay close to that prior. The router is updated under a contextual preference budget of BBB samples:

min⁡ψ1B∑i=1BLi,fψ(xi,ci)∈ΔK.\min_{\psi} \frac{1}{B} \sum_{i=1}^B \mathcal{L}_i, \qquad f_\psi(x_i,c_i)\in\Delta^K.ψmin​B1​i=1∑B​Li​,fψ​(xi​,ci​)∈ΔK.

The resulting update can be viewed as a Hedge-style online routing strategy. For each context-aware example, the method computes soft target weights over the KKK reward heads:

fi,k=ωi,kexp⁡(−Li,k/τ)∑j=1Kωi,jexp⁡(−Li,j/τ).f_{i,k} = \frac{ \omega_{i,k} \exp\left(-\mathcal{L}_{i,k}/\tau\right) }{ \sum_{j=1}^K \omega_{i,j} \exp\left(-\mathcal{L}_{i,j}/\tau\right) }.fi,k​=∑j=1K​ωi,j​exp(−Li,j​/τ)ωi,k​exp(−Li,k​/τ)​.

These soft targets emphasize reward heads that better explain the observed preference, while the temperature τ\tauτ controls sharpness of the expert weighting and the KL term prevents overly aggressive weight shifts. The router is then fine-tuned by minimizing cross-entropy between its predicted weights and the Hedge-derived soft labels:

Lrouter=−1B∑i=1B∑k=1Kfi,klog⁡fψ,k(xi,ci).\mathcal{L}_{\mathrm{router}} = -\frac{1}{B} \sum_{i=1}^B \sum_{k=1}^K f_{i,k} \log f_{\psi,k}(x_i,c_i).Lrouter​=−B1​i=1∑B​k=1∑K​fi,k​logfψ,k​(xi​,ci​).

This second-stage design keeps the routing process lightweight. The large reward functions learned during mixture modeling provide reusable preference experts, while the router adapts their mixture weights using only a small number of contextual preference examples. In deployment, the final personalized preference score is produced by combining the fixed or pre-trained reward heads with the context-aware routing weights generated by the fine-tuned router.

Experiment

The experiments evaluate a context-aware mixture reward framework on binary-labeled preference datasets, including HelpSteer2, RPR, and preference-700K, comparing it against single-reward, static mixture, ensemble adaptation, and fully supervised baselines. The results show that mixture modeling extracts specialized reward heads across diverse preference dimensions, while context-aware routing improves adaptability over fixed mixtures and related methods. In personalized settings, the framework performs well both within dataset distributions and on unseen user preferences, demonstrating label-efficient personalization with small context budgets.

MiCRo uses a mixture of Bradley-Terry rewards together with context-aware weight learning, so different prompts can activate different subpopulation reward models. Compared with prior methods that rely on fixed mixtures or no context conditioning, this enables contextual specialization and more adaptive personalization. Experiments describe it as outperforming binary-label baselines and generalizing to unseen users while remaining competitive with methods using stronger supervision. Context conditioning distinguishes the method from prior approaches that use fixed or unconditional mixture weights. Context-aware routing improves preference learning and can generalize to unseen users with a small number of contextual examples.

On the HelpSteer2 evaluation, MiCRo with context-aware routing achieves the best overall average accuracy among the compared baseline methods and improves over the single reward baseline across several attributes. The gains are most visible in helpfulness, correctness, and coherence, while some baselines remain competitive in verbosity or complexity. These results suggest that using user context for routing can support effective personalized preference adaptation with limited supervision. Context-aware MiCRo gives the strongest overall average accuracy among compared baselines. It improves helpfulness, correctness, and coherence relative to the Single Reward baseline. It outperforms the larger ARMO model and HyRe while using binary labels with context.

On the RPR test set, MiCRo achieves the best average accuracy among methods trained with binary labels, showing that context-aware routing can adapt rewards to user preferences. It remains competitive with the stronger ARMO baseline despite using fewer supervision signals, while MiCRo-HyRe gives the highest average accuracy overall. Context-aware routing improves RPR accuracy over the single reward baseline on several preference attributes. Adding test-time adaptation to MiCRo mixture heads produces the strongest overall RPR average. MiCRo offers a practical label-efficient alternative to fine-grained or test-labeled adaptation methods.

Under the unseen user setting on the RPR test set, MiCRo obtains the strongest average result among binary-label baselines trained with preference-700K. It shows notable gains on clarity, user-friendliness, storytelling, and factual accuracy, while other methods remain competitive on some dimensions such as pedagogical and linguistic creativity. Models trained on HelpSteer2 generally reach lower averages, supporting the benefit of context-aware routing for transferring across user distributions. MiCRo attains the best average score among binary-label baselines in the preference-700K group. MiCRo performs well on clarity, user-friendliness, storytelling, and factual accuracy, and ties for the top humor score. Single-head, shared-base, and static-mixture baselines remain competitive on attributes such as pedagogical quality, linguistic creativity, and scientific rigor. Context-aware routing improves adaptation to unseen users, while attribute-specific variation indicates that some preferences remain harder to learn.

Experiments evaluate MiCRo on HelpSteer2 and RPR, testing context-aware reward routing against single-reward, static-mixture, and stronger supervised baselines. The results show that context conditioning helps adapt rewards to user preferences, improving personalized preference learning and yielding strong overall performance among binary-label methods. Additional unseen-user evaluation indicates that context-aware routing transfers across user distributions while remaining competitive with methods using stronger supervision, though some preference dimensions remain harder to learn.


Build AI with AI

From idea to launch — accelerate your AI development with free AI co-coding, out-of-the-box environment and best price of GPUs.

AI Co-coding
Ready-to-use GPUs
Best Pricing

HyperAI Newsletters

Subscribe to our latest updates
We will deliver the latest updates of the week to your inbox at nine o'clock every Monday morning
Powered by MailChimp