On the Low-Rank Parametrization of Reward Models for Controlled Language Generation

TL;DR

This paper introduces a low-rank parametrization for external expert models guiding language generation, achieving comparable performance to higher-rank models with improved efficiency and reduced computational costs.

Contribution

It proposes a low-rank expert model for reward-guided decoding that maintains effectiveness while significantly reducing computational complexity.

Findings

Low-rank RAD matches high-rank RAD performance on detoxification and sentiment tasks.

Low-rank expert requires only one reward model call per token, enhancing efficiency.

Low-rank approach simplifies reward-augmented decoding without sacrificing quality.

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 an effective and modular approach for controllability of the language models, when an external expert model guides the decoding. Particularly, we zoom in into the parametrization choice of an external expert, highlighting the difference between low-rank and higher-rank parametrizations. Higher-rank experts are designed to support high flexibility when representing the rewards, leading to higher computational costs during decoding. However, we demonstrate that they might not use their full flexibility. By analyzing the recently proposed reward-augmented decoding approach (RAD), which uses a higher-rank expert model, we introduce a simpler but more efficient low-rank parametrization of the expert model enabling fast and effective guided decoding. We empirically show that the low-rank RAD performs on par with the more flexible RAD on a detoxification and a sentiment control task, while requiring only a single reward model call per generated token.