UniARM: Towards a Unified Autoregressive Reward Model for Multi-Objective Test-Time Alignment

TL;DR

UniARM introduces a unified autoregressive reward model for multi-objective test-time alignment, addressing feature entanglement and high inference costs of prior approaches by learning a single reward model with Preference-Modulated & Shared Low-Rank Adaptation (MoSLoRA). By conditioning on a mixed preference vector, UniARM jointly models all objectives in a shared parameter space and guides a frozen LLM to generate outputs aligned with user trade-offs. Training uses a combination of per-dimension and global losses to balance individual objectives and overall preferences, yielding enhanced Pareto fronts and higher alignment metrics across safety and helpfulness tasks. The approach demonstrates robust improvement over state-of-the-art methods without increasing learnable parameters or inference latency, suggesting strong practical impact for flexible, scalable multi-objective alignment in large-scale LLM deployments.

Abstract

Multi-objective alignment aims to align LLM responses with multiple human preference objectives. Among existing methods, guiding the generation of frozen LLMs through autoregressive reward models (ARMs) to accomplish multi-objective test-time alignment is a low-cost solution. However, these methods typically rely on independent parameters for each preference objective, either by training ARMs independently across preference dimensions, which neglects interactions among preference features, or by training a single ARM with separate feature extraction modules for each preference, which can cause feature entanglement. Both strategies can result in misalignment between generated outputs and user preferences. To address this limitation, we propose Preference-Modulated \& Shared Low-Rank Adaptation (MoSLoRA) for ARM training, which first extracts shared features via a preference-agnostic module and then applies affine transformations to shared features via a preference modulation module conditioned on mixed preference vectors. This design mitigates feature entanglement and enables precise control over preference trade-offs during inference. Building on this, we introduce the Unified Autoregressive Reward Model (UniARM), a novel framework for multi-objective test-time alignment. UniARM jointly models all preference dimensions in a single parameter space, eliminating the need for independent parameters for each preference objective. es on larger-scale LLMs, enhancing its practical usability.

Figures (6)

• Figure 1: Comparison between MoSLoRA and PBLoRA. MoSLoRA differs significantly in its functional design from existing methods. Unlike PBLoRA used in PARM, MoSLoRA does not learn preference-specific features. Instead, it applies affine transformations to shared features through a preference modulation module, enabling feature sharing and flexible modulation across different preferences.
• Figure 2: Comparison of the performance of various multi-objective alignment methods on safety alignment and helpful assistant tasks. w2s indicates responses generated by guiding a strong model with a weak model. Baselines RS rame2023rewarded and MOD shi2024decoding are excluded due to the high computational cost of training 13B and 65B LLMs, and their results are temporarily set to 0. Detailed experimental settings and results are provided in Section \ref{['sec:experiments']}.
• Figure 3: Learned Pareto fronts of UniARM and baseline methods on the safety alignment task. GenARM, PARM, and UniARM are fine-tuned from the Alpaca-7B model. (a) ARM guiding the generation of the frozen Alpaca-7B model. (b) ARM guiding the generation of the frozen Alpaca-65B model.
• Figure 4: Learned Pareto fronts of UniARM and baseline methods on the helpful assistant task. Figure (a) presents a 3D visualization while Figures (b), (c), and (d) show 2D projections by fixing one of the preference weights to zero. All methods are trained on the Tulu-2-7B model and then used to guide generation by the frozen Tulu-2-13B model.
• Figure 5: Learned Pareto fronts of different configurations of UniARM on the safety alignment task.