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

TL;DR

Dialogue Action Tokens (DAT) address long-horizon goal-directed dialogue by decoupling planning from generation: a small planner predicts a low-dimensional action vector, mapped via a fixed up-mapping matrix to prefix tokens that steer a frozen language model. The approach reframes dialogue as a continuous-control problem solvable with offline reinforcement learning (TD3+BC), enabling stable multi-turn planning without language deterioration. Empirically, DAT improves social capability on the Sotopia platform and reveals a potential multi-turn red-teaming vulnerability, suggesting both benefits and safety considerations of planning-based LM control. The work highlights opportunities for richer action spaces and interpretability of action-token geometry while outlining ethical implications for deployment.

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.

Figures (5)

• Figure 1: A sketch of the proposed Dialogue Action Tokens (DAT) technique. In a two-party dialogue between LM agent P (blue, part of the environment) and Q (red, DAT-steered), a multi-turn planner is introduced to steer Q towards a higher long-horizon reward. In each round of the dialogue, the planner takes the last-token embedding of the conversation history (encircled) to predict an action vector, which is then used for controlled generation ( \ref{['fig:control']}) by LM agent Q.
• Figure 2: The controlled generation process of Dialogue Action Tokens (DAT) takes two steps (sub-figure A and B). $x_t^l$ denotes the feature of the $t$-th token at the $l$-th layer. In the first step, the last-layer last-token feature is extracted as a summary of the dialogue history. Based on this summarization, the planner predicts prefix tokens li2021prefix. These prefix tokens (dialogue action tokens) are then prepended to the dialogue history tokens for controlled generation.
• Figure 3: Reward distribution in the collected replay buffer for red teaming RL training.
• Figure 4: Average success rate (ASR, $\%$) of the DAT attack on the first red teaming scenario, varying along different buffer size.
• Figure 5: Average success rate (ASR, $\%$) of the DAT attack on the first red teaming scenario, varying along number of prefix tokens ($L$), and action space dimensionality ($d'$).