Efficient Controlled Language Generation with Low-Rank Autoregressive Reward Models

2024-07-05

AI Summary

To address the challenges of uncontrollable generation and high decoding overhead in large language models (LLMs), this paper proposes Low-rank RAD, a low-rank autoregressive reward model. We first identify and theoretically characterize redundant representation in the reward matrix of standard RAD. Building on this insight, we introduce a novel lightweight reward modeling paradigm grounded in low-rank decomposition. Our method compresses reward computation to just one model call per token, achieving inference latency reduction while preserving detoxification and sentiment control performance comparable to standard RAD. By integrating task-specific reward modeling with autoregressive optimization, Low-rank RAD enables efficient, controllable, and scalable content-guided decoding. This work bridges theoretical analysis and practical deployment, offering a principled and engineering-friendly solution for controllable text generation.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work

Abstract

Language models trained on large amounts of data are known to produce inappropriate content in some cases and require careful tuning to be used in the real world. We revisit the reward augmented decoding (RAD) approach to control the generation from a language model using the scores from a task-specific reward model. We investigate the training objective of RAD, and reformulate it as a task of learning a reward matrix. We show that RAD is designed to support high flexibility when representing the reward matrices, which leads to a higher computational costs during decoding. However, we demonstrate that RAD does not use its full flexibility. Motivated by this, we propose a simpler but more efficient low-rank parametrization of the reward model enabling fast and effective guided decoding. For the detoxification and sentiment control tasks, we show that our low-rank reward model performs on par with the more flexible RAD parametrization, while requiring only a single reward model call per generated token.

Problem

• Control language models to avoid inappropriate content
• Compare low-rank vs high-rank reward model parametrizations
• Optimize reward model efficiency for guided decoding

Innovation

• Low-rank expert model for efficient decoding
• Modular approach for controlled language generation
• Single reward model call per token