Learning to Plan for Language Modeling from Unlabeled Data

By training to predict the next token in an unlabeled corpus, large language models learn to perform many tasks without any labeled data. However, their next-token-prediction objective arguably limits their performance in scenarios that require planning, such as writing a coherent article. In this paper, we train a module for planning the future writing process via a self-supervised learning objective. Given the textual context, this planning module learns to predict future abstract writing actions, which correspond to centroids in a clustered text embedding space. By conditioning on these actions, our model extends the successful language model formula to more abstract planning in an unsupervised way. Empirically, we demonstrate that our method improves language modeling performance in general, particularly with respect to the text structure. Because our framework uses a planner module that is unsupervised and external to the language model, new planner modules can be trained at large scale and easily be shared with the community.

Paper

Similar papers

Reviewer Btei7/10 · confidence 3/52024-05-14

Summary

The paper proposes to improve language models by including a planner that predicts the next writing action, in addition to the tokens. The approach is to first encode all the sentences of the model, then cluster them into writing actions, use a planner to predict the next action and fine-tune the LM to condition on the predicted planning actions. === Update after rebuttal: After reading all of the other reviews and the authors' responses, I am updating my score from 6 to 7. One thing that could be improved, is phrasing in the beginning of the paper what exactly are actions and what planning them means. At first, it is not very clear and I had the impression that you were talking about edit actions (later it was specified, but it would be good to clarify this in the beginning).

Rating

7

Confidence

3

Ethics flag

1

Reasons to accept

The proposed approach of using an action planner is meaningful, and seems to improve the model performance.

Reasons to reject

From the beginning of the paper, it is not very clear what the writing actions are, which makes it a bit hard to understand what exactly is the idea and what the model is trying to learn. The evaluation is a bit hard to follow - particularly the separation of Predicted/Fixed/Oracle.

Questions to authors

- What does "internal planner" mean in this context - is this fine-tuning of the model? - It could be interesting to discuss the complexity of the model with the added planner - does the model perform much slower, and what is the trade off between memory requirements and the improvement of the language model?

Reviewer Dmpy7/10 · confidence 4/52024-05-14

Summary

This work proposed a hierarchical way of generating text, where a high-level planning action is generated before generating the tokens in a sentence. These actions are induced by clustering sentence embeddings into K clusters, and are used to train two models: an action prediction model which predicts the action of the next sentence from the current sentence using a transformer; and a pretrained language model which is finetuned by injection actions into it. Experiments are conducted on language modeling, and the main evaluation happens at two levels: at a surface level, this method achieves a lower PPL compared to baselines; and a high-level evaluation of actions in a conditional setting, where the actions corresponding to the predicted continuation are closer to the ground truth actions compared to baselines. ===Post Rebuttal Update=== After reading authors' rebuttal, my concerns about surface-level and high-level generation evaluation have been resolved (Weakness 2 and 3). Besides, authors have clarified how the model knows sentence boundaries at generation time as well as the number of introduced parameters (which seems small). Overall, I like the idea of separating high-level planning with low-level realization, and would love to see followup works on pretraining a generally useful high-level writing planner. Therefore, I have bumped up my rating from 6 to 7.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

1. It is an interesting idea to separate high-level planning with low-level surface realization. 2. Experiments showed improved results both at low-level realization and high-level planning.

Reasons to reject

1. Some details are not clear. For example, the actions are at a sentence level, but it is not clear from the paper how the model knows when to transition to the next action at inference time: is it by introducing a special end-of-sentence token? Or by using some punctuations like "." or "?" to mark sentence boundaries? 2. The high-level action evaluation using the edit distance between the actions of ground truth continuation and generated continuation does not consider that there might be multiple possible ways to continue writing an article (even at the high level). It would be nice to see results using better evaluations, such as comparing the distribution of generated action sequences versus the original action sequences, as in "Model Criticism for Long-Form Text Generation" (https://aclanthology.org/2022.emnlp-main.815/). 3. For surface-level evaluation, this work mainly relies on PPL, but not on metrics that directly evaluate the generations themselves. It would be nice to see evaluation results under open-ended generation metrics such as MAUVE (https://openreview.net/pdf?id=Tqx7nJp7PR). 4. The method introduces additional parameters (the sentence BERT, and the transformer used in planner) compared to the baseline language model. I think it would be nice if the authors show a parameter count comparison.

Questions to authors

1. According to table 2, isn't no P init better in terms of PPL and edit distance? 2. It seems that the action prediction doesn't take into account previous actions but instead only previous text, is the rationale that previous text contains information about previous actions? 3. What's the intuition that finetuning using oracle actions is worse than using predicted actions? Is it some sort of exposure bias argument? 4. Can you visualize an example article by color-coding actions in a full article? I think that would be more useful than table 6. 5. I think this statement "However, there remain differences of 1.63 PPL and 0.47 PPL with respect to Oracle ( Oracle/Oracle ), indicating that less than half of the potential of integrating the information in the writing action is actuated due to a suboptimal planner" doesn't consider the fact that the oracle actions might leak information about the text, so the PPL it achieves might not be achievable. For example, if you use a very large action space, then each action might already be able to recover every word in the sentence, and the PPL (at least for training) would be very low.

Reviewer BMkB6/10 · confidence 4/52024-05-23

Summary

This paper presents a method for explicit planning for LLMs for text generation: 1. It breaks text generation into “chunks” (or sentences”); and for generating the next chunk, it explicitly injects an “action” variable, obtained from the planner module, to condition the next chunk generation. 2. It trains a separate transformer model to predict a sequence of action variables, as the “plan” of the generation. 3. The action variable is obtained by clustering the sentence level embeddings for sentences in the training corpus, which means no explicit supervision. Empirically, the authors find that the planning can somewhat improve generation quality (in terms of lower perplexity).

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

This paper presents an interesting idea that explicitly models the “planning” of the writing then uses it to guide/improve the generation. It automatically derives the planning in an unsupervised fashion using sentence-level embeddings of the training data. It separates the planning module and the language model, and the needed training is relatively lightweight. This idea is relevant to some of the previous work like sentence-level language modeling (Ippolito et al) and FUDGE (Yang and Klein).

Reasons to reject

This paper can be improved as follows: Methodology: 1. The reference of “plan” / “action” is based on the sentence embeddings rather than explicit writing plans like “write a sentence about concept X”, “detail the characteristics of the city”. In essence, the planning module is doing sentence level language modeling. I would recommend the authors clarify that in the paper and I would not call the $a$ variables as “actions” but just next sentence embeddings. 2. Because the author uses clusters of semantic embeddings, the granularity for expressing the fine-grained actions for the next sentence might not be fine-grained enough to guide the next generation. For example, in table 3 and table 6, the authors show examples of the major clusters of the action vectors – it only reveals major semantic/topic differences like business/geograph. That fails to capture the fine-grained control of the details during the writing process – e.g., two consecutive sentences might talk about the same general topic but focus on different aspects/details. Hence, the usefulness of the planning module is actually unclear (which is also reflected in the empirical results). 3. While it’s good to decouple the planning module from the language models, currently the training of the planning module does not use any signal from the language models, which might make the plans less helpful. Alternatively, one can use sentence representation from the same language model to train the planning module (rather than sentence bert). Experimental design and results: 1. Missing baselines: In table 1, the author only compares with a baseline that does not train the language model (the None column). At least the baseline where simply fine-tuning the language model on the same training data using a regular causal LM objective should be compared. 2. Perplexity might not be the best measurement, and one should compare the actual generation quality (In fact, improving perplexity from 27.89 to 26.69 might not necessarily indicate better generation quality). It would be helpful to directly compare generation quality, e.g., comparing rouge-2 of the generations, with/without using the plan. 3. Current table 1 evaluation (edit distance) is done at the sentence/chunk level, which might not be a realistic setup. It would be interesting to compare the quality of generating the complete article.

Questions to authors

There are some missing references: FUDGE: Controlled Text Generation With Future Discriminators For an in depth analysis of controlled generations, table 1 in Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo could be a helpful pointer.

Reviewer BMkB2024-06-06

Thanks for rebuttal!

Thank you for your responses and clarifications. Overall I like the idea of the framing of this work. Based on the rebuttal and the other reviews, I want to increase my score from 4 to 6. I am not fully convinced with the framing of the “action vector” and some parts of the experiments. I’ll add some quick thoughts below for further discussion for the future, but I do not expect the authors to respond before the rebuttal deadline. > - We do not agree that the proposed actions are called “next sentence embeddings” as they represent a cluster of sentences at a more abstract level. Consequently, the planning model does not perform “next sentence language modeling” for the same reason. > - As discussed in Section 3.1, we posit that it is necessary to compute abstractions that correspond to common writing patterns. By nature, planning in natural language involves a form of abstraction to deal with the vast action space. As Figure 2 demonstrates and we discuss in Section 5.3, there is a tradeoff between the granularity of actions (and hence informativeness for generation) and the difficulty of predicting the correct action. In the case that two consecutive sentences are guided by the same action, the language model provides fine-grained control conditioned on previous words and sentences through attention. - I partly agree with you and I would take back my original statement a bit. But I think the question still exists: starting (even clusters of) sentence representation roots the vector based on the sentence representation space. An alternative way is to purely learn this action as a latent variable model: P(a|x). That might result in a completely different representation space and still might be somewhat helpful. Another alternative is the PEER paper that specifies the editing actions: https://arxiv.org/abs/2208.11663. Regarding the “tradeoff between the granularity of actions (and hence informativeness for generation) and the difficulty of predicting the correct action”, this really boils down to a broader question: if the action is just “vague” and hard to model, then it’s unclear whether it is helpful or in which way it can help the LM. > We have tested your suggested alternative way of using a signal from the language model, namely in our ablation that encodes the entire context using GPT-2 at once rather than encoding each context sentence with a sentence transformer and then fused into one representation with a trained transformer. We observed that this yielded no improvement (see Table 2 second row and third paragraph of Section 3.2). We will state explicitly that this ablation uses GPT-2 in the final paper. - That’s good! Sorry I somehow missed that part when writing the review.

Authorsrebuttal2024-06-07

Thank you for engaging with our rebuttal!

We agree that it is preferable to have action representations that can evolve dynamically with the data. For this reason, we let action embeddings adapt when passing them to the LM. Nonetheless, we consider learning better actions from text as one of the top avenues for future work. Thank you for your suggestions in this regard.

Authorsrebuttal2024-06-07

Thank you for engaging with our rebuttal!

Given the extra page in the camera-ready version, we will add clearer definitions of actions and planning with them early in the paper.

Program Chairsdecision2024-07-10

Decision

Accept

© 2026 NYSGPT2525 LLC