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

TL;DR

UniARM introduces a unified autoregressive reward model that effectively aligns large language models with multiple human preferences at test time, reducing parameter complexity and improving preference trade-off control.

Contribution

The paper proposes UniARM, a novel framework that models multiple preferences jointly in a single parameter space, addressing limitations of previous independent or entangled approaches.

Findings

Unified modeling of preferences improves alignment accuracy.

Shared feature extraction reduces model complexity.

Enhanced control over preference trade-offs during inference.

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.