ICML 2024: NADOv2: Improved Training and Low-Rank Adaptation of Neurally-Decomposed Oracles for Controlling Language Models

Abstract

NeurAlly-Decomposed Oracle (NADO) is a powerful approach for controllable generation with large language models. It is designed to avoid catastrophic forgetting while achieving guaranteed convergence to an entropy-maximized closed-form optimal solution with reasonable modeling capacity. Despite existing success, several challenges arise when NADO is applied to less ideal scenarios. Vanilla NADO suffers from gradient vanishing for low-probability control signals and is highly reliant on the forward-consistency regularization. In addition, the vanilla implementation of NADO through introducing a few additional transformer layers suffer from a limited capacity, especially compared to other finetune-based model adaptation methods like LoRA. In this paper, we concern an improved version of the NADO algorithm, namely NADOv2, in both parameterization and training process. We discuss how such improved version can significantly improve the effectiveness of the algorithm, and allowing NADO to be combined with LoRA, achieving better model capacity and algorithmic flexibility. Experiment results on the lexically constrained generation task CommonGen justify the significance of the improvements.

TBD

Status

Accepted to ICML 2024