# SAGE: Steering Dialog Generation with Future-Aware State-Action Augmentation

arXiv: https://arxiv.org/abs/2503.03040 · HTML: https://mlsanity.com/paper/2503.03040/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 39 references that are in the map.
- First posted on arXiv 4 Mar 2025.

## Builds on (most cited first, 24 of 24)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [A Survey on Large Language Model based Autonomous Agents](https://mlsanity.com/paper/2308.11432/) (2023)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [CommonsenseQA: A Question Answering Challenge Targeting Commonsense Knowledge](https://mlsanity.com/paper/1811.00937/) (2018)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [OpenHands: An Open Platform for AI Software Developers as Generalist Agents](https://mlsanity.com/paper/2407.16741/) (2024)
- [DialoGPT: Large-Scale Generative Pre-training for Conversational Response Generation](https://mlsanity.com/paper/1911.00536/) (2019)
- [MultiWOZ -- A Large-Scale Multi-Domain Wizard-of-Oz Dataset for Task-Oriented Dialogue Modelling](https://mlsanity.com/paper/1810.00278/) (2018)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [USR: An Unsupervised and Reference Free Evaluation Metric for Dialog Generation](https://mlsanity.com/paper/2005.00456/) (2020)
- [A Hierarchical Latent Variable Encoder-Decoder Model for Generating Dialogues](https://mlsanity.com/paper/1605.06069/) (2016)
- [A Systematic Survey on Large Language Models for Algorithm Design](https://mlsanity.com/paper/2410.14716/) (2024)
- [ChiPFormer: Transferable Chip Placement via Offline Decision Transformer](https://mlsanity.com/paper/2306.14744/) (2023)
- [Zero-Shot Goal-Directed Dialogue via RL on Imagined Conversations](https://mlsanity.com/paper/2311.05584/) (2023)
- [Large Language Models for Education: A Survey](https://mlsanity.com/paper/2405.13001/) (2024)
- [MA-RLHF: Reinforcement Learning from Human Feedback with Macro Actions](https://mlsanity.com/paper/2410.02743/) (2024)

## Abstract (arXiv)

Recent advances in large language models have demonstrated impressive capabilities in task-oriented applications, yet building emotionally intelligent chatbots that can engage in natural, strategic conversations remains a challenge. We present a novel approach called SAGE that uses latent variables to control long-horizon behavior in dialogue generation. At the core of our method is the State-Action Chain (SAC), which augments standard language model fine-tuning by introducing latent variables that encapsulate emotional states and conversational strategies between dialogue turns. During inference, these variables are generated before each response, enabling coarse-grained control over dialogue progression while maintaining natural interaction patterns. We also introduce a self-improvement pipeline that leverages dialogue tree search, LLM-based reward modeling, and targeted fine-tuning to optimize conversational trajectories. Our experimental results show that models trained with this approach demonstrate improved performance in emotional intelligence metrics while maintaining strong capabilities on LLM benchmarks. The discrete nature of our latent variables facilitates search-based strategies and provides a foundation for future applications of reinforcement learning to dialogue systems, where learning can occur at the state level rather than the token level. https://github.com/apple/ml-sage-dialog-gen