DiNADO: Norm-Disentangled Neurally-Decomposed Oracles for Controlling Language Models

NeurAlly-Decomposed Oracle (NADO) is a powerful approach for controllablegeneration with large language models. It is designed to avoid catastrophicforgetting while achieving guaranteed convergence to an entropy-maximizedclosed-form optimal solution with reasonable modeling capacity. Despite thesuccess, several challenges arise when apply NADO to a wide range of scenarios.Vanilla NADO suffers from gradient vanishing for low-probability controlsignals and is highly reliant on a regularization to satisfy the stochasticversion of Bellman equation. In addition, the vanilla implementation of NADOintroduces a few additional transformer layers, suffering from a limitedcapacity especially compared to other finetune-based model adaptation methodslike LoRA. In this paper, we propose a improved version of the NADO algorithm,namely DiNADO (norm-Disentangled NeurAlly-Decomposed Oracles), which improvesthe performance of the NADO algorithm through disentangling the step-wiseglobal norm over the approximated oracle R-value for all potentialnext-tokens, allowing DiNADO to be combined with finetuning methods like LoRA.We discuss in depth how DiNADO achieves better capacity, stability andflexibility with both empirical and theoretical results. Experiments onformality control in machine translation and the lexically constrainedgeneration task CommonGen demonstrates the significance of the improvements.