# End-to-end Planner Training for Language Modeling

arXiv: https://arxiv.org/abs/2410.12492 · HTML: https://mlsanity.com/paper/2410.12492/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 42 references that are in the map.
- First posted on arXiv 16 Oct 2024.

## Builds on (most cited first, 23 of 23)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [MPNet: Masked and Permuted Pre-training for Language Understanding](https://mlsanity.com/paper/2004.09297/) (2020)
- [Probing Classifiers: Promises, Shortcomings, and Advances](https://mlsanity.com/paper/2102.12452/) (2021)
- [Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks](https://mlsanity.com/paper/1506.03099/) (2015)
- [OLMo: Accelerating the Science of Language Models](https://mlsanity.com/paper/2402.00838/) (2024)
- [Sequence Level Training with Recurrent Neural Networks](https://mlsanity.com/paper/1511.06732/) (2015)
- [Datasets: A Community Library for Natural Language Processing](https://mlsanity.com/paper/2109.02846/) (2021)
- [MAUVE: Measuring the Gap Between Neural Text and Human Text using Divergence Frontiers](https://mlsanity.com/paper/2102.01454/) (2021)
- [Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking](https://mlsanity.com/paper/2403.09629/) (2024)
- [Professor Forcing: A New Algorithm for Training Recurrent Networks](https://mlsanity.com/paper/1610.09038/) (2016)
- [A Survey of Controllable Text Generation using Transformer-based Pre-trained Language Models](https://mlsanity.com/paper/2201.05337/) (2022)
- [Hierarchical Multiscale Recurrent Neural Networks](https://mlsanity.com/paper/1609.01704/) (2016)
- [Gradient Estimation with Stochastic Softmax Tricks](https://mlsanity.com/paper/2006.08063/) (2020)
- [Guiding Language Model Reasoning with Planning Tokens](https://mlsanity.com/paper/2310.05707/) (2023)
- [Model Criticism for Long-Form Text Generation](https://mlsanity.com/paper/2210.08444/) (2022)
- [Narrative Text Generation with a Latent Discrete Plan](https://mlsanity.com/paper/2010.03272/) (2020)
- [Learning to Plan for Language Modeling from Unlabeled Data](https://mlsanity.com/paper/2404.00614/) (2024)
- [Learning Multiscale Transformer Models for Sequence Generation](https://mlsanity.com/paper/2206.09337/) (2022)

## Abstract (arXiv)

Through end-to-end training to predict the next token, LLMs have become valuable tools for various tasks. Enhancing their core training in language modeling can improve numerous downstream applications. A successful approach to enhance language modeling uses a separate planning module to predict abstract labels of future sentences and conditions the LM on these predictions. However, this method is non-differentiable, preventing joint end-to-end tuning of the planner with the LM. We propose an effective method to improve this approach by enabling joint fine-tuning of the planner and the LM. We show that a naive way of approximating the gradient of selecting a label via the straight-through estimator is not effective. Instead, we propose to use the predicted label probabilities as mixing weights to condition the LM on a weighted average of label embeddings in a differentiable manner. This not only enables joint fine-tuning of the planner and the LM, but also allows the LM to draw on the full label distribution predicted by the planner, retaining more information. Our experimental results show consistent improvements in perplexity.