Source-linked AI summary
Sequence Tutor: Conservative Fine-Tuning of Sequence Generation Models with KL-control
Natasha Jaques, Shixiang Gu, Dzmitry Bahdanau, José Miguel Hernández-Lobato, Richard E. Turner, Douglas Eck
TL;DR
Sequence generation models can produce sequences lacking global structure, while task-specific rewards alone may not preserve information learned from data or sample diversity. Sequence Tutor combines RL with an MLE model as a prior through KL-control, and demonstrations on melodies and molecules show improved structure and targeted properties while retaining data knowledge and diversity.
Problem
MLE-trained RNNs can generate sequences with poor global structure and repetition, while combining RL with task rewards may fail to preserve the data distribution and sample diversity.
Method
Sequence Tutor pre-trains an MLE RNN as a prior policy, then uses KL-control-derived off-policy RL to optimize task rewards while remaining close to that prior.
Results
Across music and molecular generation, Sequence Tutor improves desired sequence structure and targeted properties while maintaining information learned from data and sample diversity.
Takeaways & Limitations
The results provide a proof-of-concept that incomplete domain knowledge can guide RNN sequence generation without relying exclusively on imperfect rewards or data.
Takeaways & Limitations
The molecular model tends to produce simplistic, carbon-rich molecules with lower QED scores, while the music study uses simplistic monophonic representations and incomplete composition rules.
Abstract
from arXiv · showhide
This paper proposes a general method for improving the structure and quality of sequences generated by a recurrent neural network (RNN), while maintaining information originally learned from data, as well as sample diversity. An RNN is first pre-trained on data using maximum likelihood estimation (MLE), and the probability distribution over the next token in the sequence learned by this model is treated as a prior policy. Another RNN is then trained using reinforcement learning (RL) to generate higher-quality outputs that account for domain-specific incentives while retaining proximity to the prior policy of the MLE RNN. To formalize this objective, we derive novel off-policy RL methods for RNNs from KL-control. The effectiveness of the approach is demonstrated on two applications; 1) generating novel musical melodies, and 2) computational molecular generation. For both problems, we show that the proposed method improves the desired properties and structure of the generated sequences, while maintaining information learned from data.
1. Introduction
Sequence Tutor addresses the difficulty of generating globally coherent sequences by combining task-specific reinforcement learning with an MLE-trained model's data knowledge. The approach is evaluated on music generation, where it improves structure and subjective quality while reducing failure modes.
- MLE-trained sequence models can produce locally plausible text or melodies that lack coherent global structure and may repeat output tokens excessively.
- Sequence Tutor uses task-specific RL rewards while retaining information learned from data through proximity to the MLE model's distribution.
- KL-control provides off-policy RL methods that directly penalize divergence from the MLE-derived prior policy and include entropy regularization that encourages sample diversity.
- The method combines data and task-related goals without treating either the data distribution or task reward as a perfect measure of success.
- In music generation, Sequence Tutor produces melodies that are more varied, harmonious, interesting, and subjectively pleasing than those from the MLE model while reducing unwanted behaviors.
2. Related Work
Prior work combines MLE and RL for structured prediction or generative modeling, but does not address preserving data-distribution information and sample diversity in the same way. Sequence Tutor connects KL-control methods to RNN sequence generation and applies them across music and molecular domains.
- Earlier MLE-plus-RL approaches use supervised training to bootstrap RL, but RL can overwrite the original MLE policy and assume a complete task-reward specification.
- RAML augments translation targets with nearby high-reward outputs, whereas Sequence Tutor targets generation and trains an RL model rather than an MLE model.
- RAML depends on generating additional samples near ground truth that also have high reward, which may be impossible in domains where finding such samples is the problem.
- Sequence Tutor applies KL-control-inspired Ψ-learning and G-learning methods to deep RNN sequence generation and compares them with Q-learning using log-prior augmentation.
- Prior music and molecular generation studies use neural models to generate melodies or molecules, but the cited music work notes that next-note models tend to wander and lack structure.
3. Background
The background formulates reinforcement learning as sequential decision-making: an agent selects actions, receives rewards, and seeks high discounted return. Deep Q-learning approximates action values with a neural network and stabilizes training with target networks and replay.
- At each time step, an RL agent observes a state, selects an action under its policy, receives a reward, and transitions to the next state.
- The agent's objective is to maximize rewards accumulated over action sequences while discounting future rewards by γ.
- The optimal deterministic policy satisfies a Bellman optimality equation, which recursively relates action values to immediate and future rewards.
- Deep Q-learning trains a neural network to approximate Q(s, a; θ) using an objective based on rewards and the target network's next-state value.
- The target network is held fixed during gradient computation and updated more slowly, while exploration and replay memory help stabilize learning.
4. Sequence Tutor
Sequence Tutor treats the pre-trained sequence model as a prior policy and uses KL-control-based off-policy RL to improve sequence properties while retaining proximity to the data distribution. The paper derives and compares Q-learning, generalized Ψ-learning, and G-learning formulations for recurrent sequence generation.
- Objective: A trained sequence-generation RNN is treated as a black-box prior policy while RL tunes selected sequence properties without interfering with its learned distribution.The prior supplies next-token probabilities, and task-specific rewards target structure and quality.
- Objective: Figure 1 initializes the Q-network and target Q-network from an MLE-pre-trained RNN and uses a fixed copy as the Reward RNN.The Reward RNN supplies prior-policy information or cross-entropy rewards, depending on the algorithm.
- Q-learning with log prior augmentation: The simplest Q-learning variant augments task rewards with the Reward RNN output, with constant c controlling the emphasis on task-specific rewards.The resulting objective favors actions that satisfy task rules while remaining probable under the original data model.
- KL Control for Sequence Generation: Sequence generation is formulated as KL control: maximize task reward while minimizing deviation between the learned policy q(τ) and prior distribution p(τ).The sequence model is non-Markovian because each state contains all previously generated tokens.
- KL Control for Sequence Generation: The KL-control objective includes entropy regularization, so its optimal policy is stochastic rather than deterministic and can support exploration directly.This avoids relying solely on heuristic exploration methods such as ε-greedy or additive noise.
- Recurrent Ψ-learning and G-learning: The paper derives recurrent temporal-difference methods based on generalized Ψ-learning and G-learning, with G-learning representing a special case requiring dynamic mixing with prior probabilities.The methods use recurrent neural networks and experience replay; the paper compares them with Q-learning using log prior augmentation.
5. Experiment I: Melody Generation
The melody-generation experiment tests whether music-theory constraints can improve structure while retaining data-learned note probabilities. Sequence Tutor reduces targeted undesirable behaviors, preserves more information about the training distribution than RL-only training, and produces more varied and pleasing melodies, though the representation and rules remain limited.
- Experiment design: The experiment asks whether an RNN can learn music-theory constraints while retaining note probabilities learned from real songs.The rules target properties including key adherence, tonic openings, resolved leaps, harmonious intervals, motifs, and reduced repetition.
- Experiment design: The models were trained on monophonic melodies extracted from 30,000 MIDI songs and encoded as one-hot note sequences.An LSTM with one layer of 100 cells was trained before Sequence Tutor optimization.
- Rule adherence: 63% to nearly 0%: excessively repeated-note segments fell sharply after reinforcement learning across all Sequence Tutor models.The Q and Ψ models also learned to follow most imposed rules, whereas the G model improved less consistently.
- Data retention: Each Sequence Tutor model maintained higher Reward-RNN log p(a|s) than the RL-only baseline through 3,000,000 training steps.The RL-only model reached average log p(a|s) of -3.65, approximately p(a|s)=0.026, while G scored highest on this retention metric.
- Subjective quality: Q and Ψ generated the most pleasant melodies, which were more varied and interesting, stayed firmly in key, and often used harmonious intervals.The models also retained information from training data; one example ended with a seemingly familiar riff.
- Limitations: The study is a proof of concept because its monophonic encoding and simplistic, incomplete composition rules do not approach human artistic merit.The representation cannot capture musical dynamics and expressivity.
6. Experiment II: Computational Molecular Generation
Sequence Tutor is applied to molecular generation by combining an MLE-trained SMILES model with rewards for validity, drug-likeness, and synthetic accessibility. It improves most targeted metrics over the MLE baseline, but tends toward simpler, carbon-rich molecules with lower QED scores.
- Molecular generation: Only about a third of molecules generated by the MLE RNN are valid.The model generates molecules as SMILES character sequences.
- Reward design: The reward function incentivizes validity, logP, SA, and QED while penalizing carbon rings larger than six atoms.These rewards use the open-source RDkit library and incorporate a domain-specific structural constraint.
- Results: Sequence Tutor significantly improves almost all targeted molecular metrics relative to the original MLE model.The Ψ algorithm was selected for molecular generation because it performed best on the music task.
- Limitation: Sequence Tutor produces simpler, more carbon-rich molecules than MLE, lowering QED scores because simple sequences more readily achieve validity, logP, and SA rewards.The authors suggest increasing the QED reward and refining task-specific rewards with domain knowledge.
- Evaluation: Table 2 reports molecule validity and quality statistics from 100,000 randomly initialized samples, with bold entries marking significant improvements over MLE.The comparison covers validity, drug-likeness, and synthetic accessibility metrics.
7. Conclusion and Future Work
The paper concludes that Sequence Tutor uses reinforcement learning to correct sequence properties while retaining information from MLE training and preserving sample diversity. It demonstrates KL-control-based RL methods on music and molecular generation and identifies high-level domain knowledge as a broader application direction.
- Conclusion: Sequence Tutor corrects RNN-generated sequence properties while maintaining MLE-learned information and sample diversity.The framework is designed to refine, rather than discard, the data-trained behavior.
- Methodological contribution: The paper derives three RL-based sequence-generation methods through a connection between sequence generation and KL-control.These methods are empirically compared in music generation and further demonstrated for molecular generation.
- Empirical findings: Experiments show promising adherence to task-specific rules and subjective quality in generated sequences.The evidence spans music generation and computational molecular generation.
- Future work: Sequence Tutor provides an alternative way to encode high-level domain knowledge into an RNN instead of relying solely on potentially biased data.The authors identify bias reduction as one possible application.