Dialogue Action Tokens: Steering Language Models in Goal-Directed Dialogue with a Multi-Turn Planner

Authors: Kenneth Li, Yiming Wang, Fernanda Viégas, Martin Wattenberg
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: large language model, dialogue systems, social intelligence, red teaming

Abstract:
We present an approach called Dialogue Action Tokens (DAT) that adapts language model agents to plan goal-directed dialogues. The core idea is to treat each utterance as an action, thereby converting dialogues into games where existing approaches such as reinforcement learning can be applied. Specifically, we freeze a pretrained language model and train a small planner model that predicts a continuous action vector, used for controlled generation in each round. This design avoids the problem of language degradation under reward optimization. When evaluated on the Sotopia platform for social simulations, the DAT-steered LLaMA model surpasses GPT-4's performance. We also apply DAT to steer an attacker language model in a novel multi-turn red-teaming setting, revealing a potential new attack surface.

Decision: reject
Reviewer scores: 5, 6, 8, 8, 5

Reviews and discussion:
<thread>
creation date:
1729562943318

thread_type:
official_review

<thread_object>
creator: reviewer_rkaj

<content>
summary:
This work proposes a light-weight RL-based technique to improve language model's goal completion ability. The authors propose Dialogue Action Tokens (DAT), which are prefix tokens emitted by a small (trained) planner model to steer a larger LM's response at test time. During inference time, DAT uses the feature vector from the context's last word from the larger frozen LM, feeds them into a small MLP model (the planner) to output L-tokens, and finally prefixes the input texts with these L-tokens and prompt the frozen LLM for generation. During training, only the small planner model is trained using policy gradient. The author then evaluated DAT in a social intelligence benchmark (Sotopia) and a red-teaming benchmark (HarmBench), and showed improved performance compared to alternatives such as no training.

strengths:
- Since DAT only optimizes a small planner model, it is a compute-efficient approach to improve performance of LLMs.
- The authors presented diverse experiments in social intelligence benchmarks (Sotopia) as well as red-teaming benchmarks (HarmBench) to validate the effectiveness of their approach

weaknesses:
1. I believe there is a significant amount of related work being overlooked. This work proposes DAT motivated by the lack of RL techniques in optimizing utterance-level MDP (L31-39, L128-131) as well as planning/optimizing in continuous action space (L47-49), which is not true. Optimizing utterance-level MDP has been extensively explored by work including but not limited to NLPO [1], RvS [2], ILQL [3], and more. There is little to no mention to these methods, and there is no comparisons to direct RL training approaches in the experiments. Additionally, planning/optimizing in continuous vector space has also been explored by many prior work, especially in task oriented dialogues. For example, LAVA [4] and TCUP [5] optimizes the LM in a latent action space by formulating the auto-regressive generation process as variational inference, and uses RL algorithms such as policy gradient for direct optimization. These were also of high relevance but not mentioned or compared against.

2. Some design choices about DAT seems questionable. For example, the authors only uses a very small planner model, and only L=2 prefix tokens is used to steer the LLM output (presumably to prevent language degradation mentioned in the introduction). However, this presents a significant trade-off between performance v.s. robustness, as increasing L may again cause language degradation. Has the authors experimented with L>2?

3. Comparisons made in the experiments in Section 7 Sotopia were not fair. The author compared a DAT-trained model against no-training/directly prompting LLaMA-2, LLaMA-3, or GPT-4. However, DAT should be compared against *many existing RL training methods including [1-5] and even PPO [6]*, since Sotopia tasks can be formulated as an MDP problem.

4. Comparisons made in the experiments in Section 8 Red Teaming were not fair. 1) Baselines include GDG and PAIR, which constructs an adversarial prompt by searching prefixes/prompts *at test time*, whereas DAT additionally *trained the planner against the defender* before testing. Similar to section 7, there should be comparisons against direct RL methods. 2) DAT performance became *significantly worse against a much weaker defender (LLaMA-2-7b-chat)*, whereas other baselines show improved performance. Is this an indication of reward over-optimization of DAT-training?

---

References:

[1] Ramamurthy, Rajkumar et al. “Is Reinforcement Learning (Not) for Natural Language Processing?: Benchmarks, Baselines, and Building Blocks for Natural Language Policy Optimization.” ArXiv abs/2210.01241 (2022): n. pag.

[2] Emmons, Scott et al. “RvS: What is Essential for Offline RL via Supervised Learning?” ArXiv abs/2112.10751 (2021): n. pag.

[3] Snell, Charles Burton et al. “Offline RL for Natural Language Generation with Implicit Language Q Learning.” ArXiv abs/2206.11871 (2022): n. pag.

[4] Lubis, Nurul et al. “LAVA: Latent Action Spaces via Variational Auto-encoding for Dialogue Policy Optimization.” ArXiv abs/2011.09378 (2020): n. pag.

[5] Vlastelica, Marin et al. “Taming Continuous Posteriors for Latent Variational Dialogue Policies.” AAAI Conference on Artificial Intelligence (2022).

[6] Schulman, John, et al. "Proximal policy optimization algorithms." arXiv preprint arXiv:1707.06347 (2017).

questions:
In general, I believe this paper is poorly structured. For example,

- Many empirical details such as "Notes on architecture" (L168-173); "Remarks" (L235-237) should not be mentioned in the method section but rather in experimental setups.
- Experiments section should be sectioned into separate sections that introduces the benchmark, baselines, specific implementation details, results, and ablations/analysis. Currently, there are only two subsections in Section 6 and Section 7 that discusses all the relevant information mentioned above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729680085129

thread_type:
official_review

<thread_object>
creator: reviewer_sfxb

<content>
summary:
This paper presents a novel approach, Dialogue Action Tokens (DAT), to improve the long-term performance of Large Language Models (LLMs) in goal-oriented dialogues. The proposed method introduces a planner module and an up-mapping matrix to backbone LLMs, enabling explicit prediction of dialogue actions at each turn. These predicted action embeddings are prepended to the dialogue history when generating responses. Both modules can be trained in continuous vector space using Reinforcement Learning (RL) with a dedicated reward function. The training process consists of two steps: (1) cloning the original backbone LLM's behavior, which serves as pretraining for the two modules; and (2) fine-tuning the planner module with the reward function while freezing the up-mapping matrix and backbone LLM. The method is evaluated on two tasks, Social Capability and Red Teaming, demonstrating significant improvements over strong baselines.

strengths:
1. The idea of explicitly prepending dialogue action tokens to the dialogue history is novel and interesting, offering a valuable contribution to future research in the community.
2. Instead of relying on prompt engineering, the proposed method introduces lightweight, trainable parameters to LLMs, which can be trained through RL. This approach makes the training process more mathematically sound and effective.
3. The proposed method achieves impressive gains on both tasks compared to state-of-the-art results, particularly on the Red Teaming task. The results are convincing, supported by extensive experiments and ablation studies.
4. The method has good theoretical generalizability, potentially applicable not only to goal-oriented dialogue but also to non-goal-oriented scenarios if the underlying purpose can be defined. Additionally, the technique is not limited to specific backbone models or RL algorithms.

weaknesses:
1. A notable limitation is that the proposed method requires a reward function for RL, which can be hard or costly to acquire as mentioned in the paper. This may restrict the use of DAT in certain scenarios where reward signals are difficult to obtain.
2. While the evaluation results are promising, they rely solely on automatic methods. Incorporating human evaluation would provide additional validation and make the results even more convincing.

questions:
1. Are predicted action tokens at each turn also added to the dialogue history?  
2. In Table 3,  why does Llama-2-7B have better performance than Llama-3-8B agaist two of single-turn attackers?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730797794202

thread_type:
official_review

<thread_object>
creator: reviewer_lbpv

<content>
summary:
Whilst large language models have illustrated impressive performance in many tasks, it has also become clear that they often struggle to plan sufficiently when faced with longer term or challenging goals. This work proposes the use of simplified dialogue action tokens used to steer the LLM in the right direction. They use a separate, small, planning module which assist the LLM in planning for longer term goals. Through RL, this planner learns to successfully steer the LLM through a complex goal, such as those present in the Sotopia scenarios.

strengths:
- This work presents a relevant method for improving the long term planning of large language models using dialogue action tokens.
- The method is well tested in the two scenarios, Sotopia and the Red Teaming Scenarios.
- A small, and hence "cheap", planning module to insert tokens for planning improves planning performance significantly. Hence, providing a computationally efficient method to improve the performance of LLMs in long term tasks without expensive LLM fine-tuning.

weaknesses:
Whilst the small planning module does successfully steer the LLM, it does not improve the actual general planning ability of the LLM. Further, through the bottleneck of the DAT tokens, information can be lost which could have been beneficial for generating better responses. This leaves whether the LLM could learn to perform this planning unanswered.

questions:
Have you tested whether RL to generate such planning tokens directly by the LLM can be done? Possibly using a penalty to avoid diverging from natural language.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730835872037

thread_type:
official_review

<thread_object>
creator: reviewer_jzcm

<content>
summary:
The paper looks at training an LLM to perform the role of a conversational agent which has certain goals, e.g. steering a user towards a certain response. The paper approaches this by training a small module which produces essentialy control tokens (2 in the paper) that are the same dimensions as the LLM token embeddings. The LLM is the conditioned on these extra tokens when producing a response. The paper gives a good description of RL and past work, and motivates taking this approach primarily as it doesn't change the LLM, so by design doesn't degrade its core language abilities. 
This has been done before under terms like prompt learning, with differences mainly in how the "control tokens" are obtained.

strengths:
Clear motivation, tested on open data and compared against prior methods. Results indicate that the proposed method is working, although there are a few asterisks on the evaluation

weaknesses:
* The evaluation is based on the same signal being optimised by the reported model. This is a bias towards the proposed approach being better than the others. Some what tricky to avoid, but human evaluation would be one (admittedly expensive) approach to more robust evaluation. 
* The method relies on a very powerful evaluator. This is ok, but important to note (as the paper does). However there is no reporting here of how accurate even the prompted GPT was on this task. I presume such was measured as part of prompt engineering your way to a good LLM judge? It would be helpful to know how accurate the evaluator was, both for interpretation of the actual reported results, but also for how much signal (versus noise) was being input to the RL training. 
* are any of those results statistically significant? Looking at table 1 alone, I'm not sure they would be, as most seem to be overlapping when considering the 68% intervals

questions:
* Interested in the authors speculation on why the certain words were obtained as interpretations of the policy model (lines 433-434)? 
* Also, the policy only outputs 2 vectors, are all those words tokenised into 2 tokens? 
* Should the partner model be a different LLM? Presumably it is significantly easier to learn how to bias "your twin" so to speak, compared to another model? Maybe that's a speculative question worthy of future analysis. Interested in your thoughts.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730868255809

thread_type:
official_review

<thread_object>
creator: reviewer_diqr

<content>
summary:
The authors propose an approach called Dialogue Action Tokens which serves as a continuous vector representing actions that an LLM should take throughout the course of a conversation. The proposed approach is evaluated on SOTOPIA, as well as a novel task called ``multi-turn red-teaming'' designed by the authors.

strengths:
Overall, the writing is clear and the problem formulation is sound -- it is reasonable to consider as an MDP as described in Section 3. 

The general idea of adapting LLMs to downstream conversational scenarios is also important, and the generalizability of a continuous action vector is also an important concept. 

The proposed multi-turn red-teaming task is also interesting.

weaknesses:
My primary concerns with this work are that the baselines are not properly set, and the overall novelty of the proposed work is limited. The work makes a strong claim that LLMs can only be steered for downstream applications via prompt engineering, and that the proposed DAT method addresses the gap by virtue of being an RL-based approach which plans actions across long horizons. However, the work does not engage with the existing literature on multi-turn conversations. In particular, the main experimental results in Table 1/3 do not include comparisons to other dialogue action planning baselines (e.g., [1, 2, 3]). It is also not clear how the proposed work is fundamentally different than existing work looking at continuous representations of dialogue acts (e.g. [4]).

[1] Plug-and-Play Policy Planner for Large Language Model Powered Dialogue Agents, ICLR 2024

[2] Prompt-Based Monte-Carlo Tree Search for Goal-Oriented Dialogue Policy Planning, EMNLP 2023

[3] Planning Like Human: A Dual-process Framework for Dialogue Planning, ACL 2024

[4] DiactTOD: Learning Generalizable Latent Dialogue Acts for Controllable Task-Oriented Dialogue Systems, SIGDIAL 2023

questions:
Why is the reasoning behind proposing red-teaming as a multi-turn dialogue task?

Given that prompt engineering is an effective way to adapt an LLM for downstream conversational tasks (L31), have you compared DAT to any prompting approaches for goal-oriented dialogue? Can DAT be combined with any pre-existing prompting approaches? 

This isn't a question, but - I think it would be helpful to examine the literature on mixed-initiative conversational agents, if the authors intend to frame the overall work as a contribution to work on dialogue.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces Dialogue Action Tokens (DAT), a lightweight RL-based method designed to improve the goal-directed capabilities of LLMs in dialogue tasks. By using a small planner module to generate action tokens, which are prepended to the dialogue input, the approach steers the LLM without altering its core parameters, preserving its foundational abilities. Evaluated on Sotopia and a novel multi-turn red-teaming task, DAT demonstrates enhanced performance over strong baselines, particularly in long-term planning and complex goal-oriented scenarios.

Strengths:

1. The approach is compute-efficient, as it optimizes only a small planner model to improve LLM performance.
2. The multi-turn red-teaming task is interesting and effectively demonstrates the method's applicability.
3. The experiments are diverse, spanning both social intelligence benchmarks (Sotopia) and red-teaming benchmarks (HarmBench).
4. The method employs trainable parameters prepended as tokens, providing a practical alternative to prompt engineering and expensive fine-tuning, though two reviewers have concerns about the novelty and impact of this work.

Weaknesses:

1. Novelty: The submission lacks sufficient comparison with existing literature on multi-turn conversations and related long-horizon planning methods, raising concerns about the originality of the proposed approach. Furthermore, the authors do not clarify how their method fundamentally differs from existing work on continuous representations of dialogue acts. The submission also overlooks a substantial body of prior research on optimizing utterance-level MDPs and planning in continuous vector spaces, which undermines the claims of novelty. Specifically, RL-based methods such as NLPO, RvS, and ILQL have already addressed utterance-level MDP optimization, while approaches like LAVA and TCUP have explored continuous action spaces in task-oriented dialogues.

2. Experiments: The experiments do not include comparisons to relevant dialogue action planning baselines, which weakens the evaluation of the proposed method. Additionally, the comparisons in Section 7 are unfair, as the proposed method is evaluated against prompting-based baselines rather than RL-based methods, which are more appropriate for tasks formulated as MDP problems.

3. Evaluation: The reliance on automatic evaluation introduces potential biases, especially since the evaluation signal directly aligns with the model's optimization objective. Human evaluations, while costly, would add robustness to the reported results and help address concerns about evaluator accuracy and statistical significance.

4. Design choices: The paper lacks sufficient justification for certain design decisions, such as the use of a small planning module and limiting the prefix tokens to L=2. These choices appear to trade off performance for robustness, and the implications of increasing L remain unexplored.

During the discussion, the authors made significant efforts to address the reviewers' concerns and added new experiments that addressed some of the issues with the experiments (e.g., the addition of NLPO). However, the extensive discussions also revealed that the authors misunderstood key points about related work (L31-39, L128-131, and L47-49), and significant concerns about missing related work still remain. Due to the substantial amount of relevant approaches being overlooked, two reviewers and I believe the paper would require significant rewriting and reframing, which might also affect the paper's claims of novetly.

Another issue is that the authors neither updated the paper nor provided a unified response to all the reviewers explaining how they intended to address the major novelty concerns raised. This makes it difficult for me to recommend acceptance, as I would have preferred to read a unified explanation from the authors—ideally incorporated into updated introduction and related work sections—clarifying why the paper remains significant despite the related work overlooked in the original and only submission. Instead, I am left to piece together, from the various discussions, (1) the authors’ updated views on the paper's contribution and (2) their overall plan to address its shortcomings. This approach provides only a vague and incomplete picture of what we can reasonably expect in the camera-ready paper. Piecing this together is challenging because, e.g., the responses to novelty concerns were directed to two reviewers who had related yet distinct concerns.

Given all these uncertainties regarding the merits of the paper and how its shortcomings might be addressed, I recommend rejecting it.