# Controlling Large Language Model with Latent Actions

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 61 of its 69 references that are in the map.
- First posted on arXiv 27 Mar 2025.

## Builds on (most cited first, 30 of 61)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [A Survey on Hallucination in Large Language Models: Principles, Taxonomy, Challenges, and Open Questions](https://mlsanity.com/paper/2311.05232/) (2023)
- [ALFWorld: Aligning Text and Embodied Environments for Interactive Learning](https://mlsanity.com/paper/2010.03768/) (2020)
- [Qwen2.5-Math Technical Report: Toward Mathematical Expert Model via Self-Improvement](https://mlsanity.com/paper/2409.12122/) (2024)
- [M3-Embedding: Multi-Linguality, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation](https://mlsanity.com/paper/2402.03216/) (2024)
- [GLM: General Language Model Pretraining with Autoregressive Blank Infilling](https://mlsanity.com/paper/2103.10360/) (2021)
- [CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers](https://mlsanity.com/paper/2205.15868/) (2022)
- [DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs](https://mlsanity.com/paper/1903.00161/) (2019)
- [Length-Controlled AlpacaEval: A Simple Way to Debias Automatic Evaluators](https://mlsanity.com/paper/2404.04475/) (2024)
- [StarCoder: may the source be with you!](https://mlsanity.com/paper/2305.06161/) (2023)
- [Genie: Generative Interactive Environments](https://mlsanity.com/paper/2402.15391/) (2024)
- [Safe RLHF: Safe Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2310.12773/) (2023)
- [Llemma: An Open Language Model For Mathematics](https://mlsanity.com/paper/2310.10631/) (2023)
- [Cognitive Behaviors that Enable Self-Improving Reasoners, or, Four Habits of Highly Effective STaRs](https://mlsanity.com/paper/2503.01307/) (2025)
- [Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos](https://mlsanity.com/paper/2206.11795/) (2022)
- [ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search](https://mlsanity.com/paper/2406.03816/) (2024)

## Built on by (most cited first, 5 of 5)

- [A Survey on Large Language Models for Mathematical Reasoning](https://mlsanity.com/paper/2506.08446/) (2025)
- [Base Models Know How to Reason, Thinking Models Learn When](https://mlsanity.com/paper/2510.07364/) (2025)
- [Unleashing Embodied Task Planning Ability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2506.23127/) (2025)
- [Stop-Think-AutoRegress: Language Modeling with Latent Diffusion Planning](https://mlsanity.com/paper/2602.20528/) (2026)
- [Modularized Reinforcement Learning on LLMs: From MDP Creation to Exploration and Learning](https://mlsanity.com/paper/2606.21943/) (2026)

## Abstract (arXiv)

Adapting Large Language Models (LLMs) to downstream tasks using Reinforcement Learning (RL) has proven to be an effective approach. However, LLMs do not inherently define the structure of an agent for RL training, particularly in terms of defining the action space. This paper studies learning a compact latent action space to enhance the controllability and exploration of RL for LLMs. We propose Controlling Large Language Models with Latent Actions (CoLA), a framework that integrates a latent action space into pre-trained LLMs. We apply CoLA to the Llama-3.1-8B model. Our experiments demonstrate that, compared to RL with token-level actions, CoLA's latent action enables greater semantic diversity in text generation. For enhancing downstream tasks, we show that CoLA with RL achieves a score of 42.4 on the math500 benchmark, surpassing the baseline score of 38.2, and reaches 68.2 when augmented with a Monte Carlo Tree Search variant. Furthermore, CoLA with RL consistently improves performance on agent-based tasks without degrading the pre-trained LLM's capabilities, unlike the baseline. Finally, CoLA reduces computation time by half in tasks involving enhanced thinking prompts for LLMs by RL. These results highlight CoLA's potential to advance RL-based adaptation of LLMs for downstream applications.