End-to-end Planner Training for Language Modeling
Fuente:
arXiv
Saved in:
| Main Authors: | , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866917804975849472 |
|---|---|
| author | Cornille, Nathan Mai, Florian Sun, Jingyuan Moens, Marie-Francine |
| author_facet | Cornille, Nathan Mai, Florian Sun, Jingyuan Moens, Marie-Francine |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2410_12492 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | End-to-end Planner Training for Language Modeling Cornille, Nathan Mai, Florian Sun, Jingyuan Moens, Marie-Francine Computation and Language Machine Learning 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. |
| title | End-to-end Planner Training for Language Modeling |
| topic | Computation and Language Machine Learning |
| url | https://arxiv.org/abs/2410.12492 |