Controlling Large Language Model with Latent Actions

📄 arXiv: 2503.21383v1 📥 PDF

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

分类: cs.CL, cs.LG

发布日期: 2025-03-27

💡 一句话要点

CoLA：通过学习紧凑潜在动作空间，提升大型语言模型在强化学习中的可控性和探索能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 强化学习 潜在动作空间 文本生成 可控性 探索 下游任务 Llama-3

📋 核心要点

  1. 现有方法在利用强化学习适配大型语言模型时，缺乏对Agent动作空间的有效定义，限制了模型的可控性和探索能力。
  2. CoLA框架通过学习紧凑的潜在动作空间，并将其集成到预训练的LLM中，从而提升模型在强化学习中的性能。
  3. 实验表明，CoLA在数学问题求解和Agent任务中均优于基线方法，且能减少计算时间，同时保持LLM的预训练能力。

📝 摘要（中文）

本文提出了一种名为CoLA（Controlling Large Language Models with Latent Actions）的框架，旨在通过学习紧凑的潜在动作空间来增强大型语言模型（LLMs）在强化学习（RL）中的可控性和探索能力。CoLA将潜在动作空间集成到预训练的LLMs中，并应用于Llama-3.1-8B模型。实验结果表明，与使用token级别动作的RL相比，CoLA的潜在动作能够生成更多样化的语义文本。在下游任务增强方面，CoLA与RL结合在math500基准测试中达到了42.4分，超过了基线38.2分，并结合蒙特卡洛树搜索变体达到了68.2分。此外，CoLA与RL结合能够持续提升基于Agent的任务性能，且不会降低预训练LLM的能力，这与基线方法不同。最后，CoLA在涉及增强LLMs思考提示的任务中，将计算时间减少了一半。这些结果突显了CoLA在推进基于RL的LLMs下游应用适配方面的潜力。

🔬 方法详解

问题定义：现有方法在使用强化学习（RL）微调大型语言模型（LLMs）时，通常直接使用token级别的动作空间。这种方式存在两个主要问题：一是动作空间过于庞大，导致探索效率低下；二是缺乏对动作语义的有效控制，难以引导LLM生成期望的文本。因此，需要一种更紧凑、更具语义信息的动作空间来提升LLM在RL中的可控性和探索能力。

核心思路：CoLA的核心思路是学习一个低维的潜在动作空间，并将该空间映射到LLM的输出空间。通过在潜在空间中进行探索，可以更高效地学习策略，并生成更具语义一致性的文本。这种方法类似于在连续控制任务中使用潜在动作空间，但应用于离散的文本生成任务。

技术框架：CoLA框架包含以下几个主要模块：1) 潜在动作编码器：将LLM的隐藏状态编码为潜在动作向量。2) 潜在动作解码器：将潜在动作向量解码为LLM的输入，影响LLM的文本生成。3) 强化学习模块：使用RL算法（如PPO）训练潜在动作编码器和解码器，以最大化奖励函数。整个训练过程是端到端的，LLM的参数可以固定或微调。

关键创新：CoLA最重要的创新在于将潜在动作空间的概念引入到LLM的强化学习微调中。与直接使用token级别的动作空间相比，潜在动作空间更加紧凑，更易于探索，并且能够更好地控制生成文本的语义。此外，CoLA框架能够有效地平衡探索与利用，避免了LLM在微调过程中遗忘预训练知识。

关键设计：CoLA的关键设计包括：1) 潜在动作空间的维度：需要根据任务的复杂度和计算资源进行调整。2) 潜在动作编码器和解码器的网络结构：可以使用Transformer或其他神经网络结构。3) 奖励函数的设计：需要根据具体的下游任务进行定制，以引导LLM生成期望的文本。4) 强化学习算法的选择：可以使用PPO、SAC等常见的RL算法。

🖼️ 关键图片

img_0

📊 实验亮点

CoLA在math500基准测试中取得了显著的性能提升，达到了42.4分，超过了基线38.2分。结合蒙特卡洛树搜索变体后，更是达到了68.2分。此外，CoLA在Agent任务中也表现出色，能够持续提升性能，且不会降低LLM的预训练能力。更重要的是，CoLA在涉及增强LLMs思考提示的任务中，将计算时间减少了一半，这表明CoLA具有很高的效率。

🎯 应用场景

CoLA框架具有广泛的应用前景，可以应用于各种需要通过强化学习微调LLM的场景，例如：对话生成、文本摘要、机器翻译、代码生成、游戏AI等。通过学习紧凑的潜在动作空间，CoLA可以提升LLM在这些任务中的性能和可控性，并降低计算成本。未来，CoLA还可以与其他技术（如模仿学习、对抗学习）相结合，进一步提升LLM的性能。

📄 摘要（原文）

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.