Controlling Large Language Model with Latent Action

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 Controlling Large Language Models with Latent Actions 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}.

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.