Controlling Large Language Model with Latent Action

Chengxing Jia, Ziniu Li, Pengyuan Wang, Yi-Chen Li, Zhenyu Hou, Yuxiao Dong, Yang Yu

TL;DR: Develop a new architecture of language model to control language with latent action

Abstract: 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 specifying the action space. This paper studies learning a compact latent action space to enhance the controllability and exploration of RL for LLMs. Inspired by reinforcement learning from observations, we propose **Co**ntrolling Large Language Models with **L**atent **A**ctions **CoLA**, a framework that integrates a latent action space into pre-trained LLMs. **CoLA** employs an \emph{inverse dynamics model} to extract latent actions conditioned on future tokens, ensuring that the next token prediction is partially influenced by these actions. Simultaneously, **CoLA** fine-tunes the pre-trained LLM to function as a \emph{language world model}, capable of incorporating latent actions as inputs. Additionally, **CoLA** trains a \emph{policy model} to generate actions within this language world model. The policy model can be trained via behavior cloning to mimic a standard language model or through RL to maximize task-specific rewards. In this work, we apply **CoLA** to the Llama-3.1-8B model. Our experiments demonstrate that, compared to RL with token-level actions, **CoLA**'s latent actions enable greater semantic diversity. For enhancing downstream tasks, we show that **CoLA** with RL achieves a score of 42.4 on the \emph{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 via RL. These results highlight **CoLA**'s potential to advance RL-based adaptation of LLMs for downstream applications. The CoLA model is available at \url{https://huggingface.co/LAMDA-RL/Llama-3.1-CoLA-10B}.

Lay Summary:

This paper introduces a new method called **CoLA** to improve how large language models (like chatbots) can be fine-tuned for specific tasks using reinforcement learning (RL)—a technique that helps AI learn by trial and error. Normally, adjusting these models with RL can be slow and inefficient because the possible actions (like choosing words) are too broad and hard to control.

**CoLA** simplifies this process by learning a smaller, more meaningful set of “latent actions" that guide the model’s responses instead of dealing with every possible word choice. This makes the model easier to control and helps it explore better solutions faster.

In tests, **CoLA** improved performance on reasoning and agentic tasks, achieving higher scores than standard methods while also being twice as fast in some cases. Importantly, it fine-tunes the model without harming its original abilities, unlike some existing approaches. This makes **CoLA** a promising tool for adapting large language models to real-world applications more efficiently.

The model is publicly available for researchers and developers to use.

Link To Code: https://github.com/LAMDA-RL/CoLA

Primary Area: Deep Learning->Large Language Models

Keywords: Controllable language model, latent action model

Submission Number: 16114