UniARM: Towards a Unified Autoregressive Reward Model for Multi-Objective Test-Time Alignment
This work addresses alignment challenges for LLM developers and users by providing a more efficient and effective method for multi-objective test-time alignment, though it is incremental as it builds on existing autoregressive reward model approaches.
The paper tackles the problem of multi-objective alignment in LLMs by proposing UniARM, a unified autoregressive reward model that uses a novel training method (MoSLoRA) to mitigate feature entanglement and enable precise preference control, achieving competitive performance on benchmarks like AlpacaEval and MT-Bench with up to 10% improvement in alignment scores.
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.