Context-Aware Language Modeling for Goal-Oriented Dialogue Systems

TL;DR

CALM reframes goal-oriented dialogue as a POMDP and trains end-to-end language models offline to maximize task reward. It introduces task relabeling, a task-aware auxiliary loss, task pretraining, and model-based dialogue rollouts to steer generation toward task success without sacrificing language quality. On AirDialogue’s flight-booking tasks, CALM outperforms prior state-of-the-art by about 7% and reaches human-level performance under standard evaluation, while maintaining strong language metrics. This end-to-end, context-aware approach reduces reliance on hand-designed dialogue management and enables scalable, offline training with planning capabilities for practical deployment.

Abstract

Goal-oriented dialogue systems face a trade-off between fluent language generation and task-specific control. While supervised learning with large language models is capable of producing realistic text, how to steer such responses towards completing a specific task without sacrificing language quality remains an open question. In this work, we formulate goal-oriented dialogue as a partially observed Markov decision process, interpreting the language model as a representation of both the dynamics and the policy. This view allows us to extend techniques from learning-based control, such as task relabeling, to derive a simple and effective method to finetune language models in a goal-aware way, leading to significantly improved task performance. We additionally introduce a number of training strategies that serve to better focus the model on the task at hand. We evaluate our method, Context-Aware Language Models (CALM), on a practical flight-booking task using AirDialogue. Empirically, CALM outperforms the state-of-the-art method by 7% in terms of task success, matching human-level task performance.

Figures (7)

• Figure 1: CALM is an end-to-end language model for goal oriented dialogue. CALM's training objective teaches the model to better pay attention to the dialogue task context, yielding a $\sim$50% improvement in task success over standard LM training on a flight booking task.
• Figure 2: A visual outline of CALM. We apply Task Relabeling to our static offline dataset, by swapping out the task context –- in this case a flight table –- such that the attached dialogue becomes an optimal example of task completion. When fine-tuning on this relabeled dataset, we then apply a Task Specific Auxiliary Loss on top of the standard language modeling objective; this helps the model learn to use the task context. Once trained, CALM can consistently solve goal-directed dialogue tasks.
• Figure 3: Example dialogues generated by CALM (in green) in the simulated evaluation. Despite being end-to-end, CALM produces highly coherent and sensible outputs.
• Figure 4: Task success as a function of the number of rollout samples. Note that successful task completion improves with more rollout samples.
• Figure 5: Our dialogue rollout planning procedure. To generate our response, we sample entire dialogues from the language model and then re-rank the predicted dialogues with a reward function.