Controlling Large Language Model with Latent Action

We sincerely appreciate the reviewer's valuable feedback and recognition of our latent action approach.

Q1: Claims are not supported by clear evidence: Therefore, LAMP aims to decouple these semantics by assigning them to different and limited latent actions, enabling more efficient fine-tuning, alignment, and search by latent action control. ==> I did not see such evidence.

A1: For claim 1, we realize that this claim may potentially lead to misunderstandings and would like to clarify. Figure 12 and Figure 10 (Appendix C.2) aim to show that latent actions have higher semantics than tokens. We will revise the statement in Section 4 to clarify that the connection between latent actions and semantics is weakly correlated, which should avoid misunderstandings.

Q2: Claims are not supported by clear evidence: However, in our experiments, we observed that this suffers from codebook collapse, where only a limited number of actions are activated during training.

A2: For claim 2, we compare VQ-VAE with our proposed method in Figure 8 (Appendix C.1) by counting the number of actions with non-zero usage. Results show that VQ-VAE suffers from severe codebook collapse, while our approach effectively maintains the utilization of codes in the codebook. We will add a reference to this claim to make it more supported.

Q3: However, the authors should show that their methods are appliable to other LLMs besides Llama-3.1-8B.

A3: We have implemented the LAMP architecture with Qwen-2.5-math-1.5B (Appendix C.1, Figure 9). We observed similar behavioral patterns to those in Llama-3.1-8B, including stable loss reduction and maintained codebook diversity without collapse. This provides preliminary evidence for LAMP's generalizability. However, due to the limitation of computation resources, we did not proceed to full-scale training and subsequent downstream task optimization.

Q4: In the generation process, we compute the next token from the world model in a greedy manner. We only randomly or greedily sample actions from the policy model. ==> what's your meaning?

A4: This sentence describes LAMP's language generation process. During the RL training phase, LAMP's policy model explores by randomly sampling latent actions, while during the evaluation phase, it deterministically (greedily) samples latent actions to ensure consistent results by selecting the maximum probability action. We will further revise here to make a clearer description.

Q5: Some writing issues, including typos and redundant descriptions.

A5: We will improve the writing in the revision.

Q6: line-199: why the current time embedding e^i_t has a superscript i?

A6: Here we aim to differentiate the embedding from other model input/output embeddings, and $i$ represents the inverse dynamics model.

Q7: I did not see the necessary of MCTS-Q? It makes the methods complex?

A7: MCTS-Q demonstrates how our latent action formulation can effectively integrate techniques that were previously difficult to apply in token space. For instance, Q-learning-based uncertainty estimation is a key technique for pruning the search space in MCTS, but scaling it in token space is challenging. Thanks to our latent action space design, however, this approach now shows promising results. Importantly, MCTS-Q does not introduce significant complexity—instead, it makes MCTS more effective.

Q8: The results did not show that the proposed LAMP are better than baselines with a clear margin.

A8: We guess the reviewer draws this conclusion from the results in Table 1 and Table 2. We want to clarify the performance gains of our approach. Our method primarily demonstrates superior downstream performance in multiple tasks. For preference alignment, we achieve an average win rate of 64% across four tasks and exceed 90% under the KL=0.00 setting. For math reasoning, we also shows significant improvement (42.6 vs 38.2). For search, combined MCTS with Q, our method can achieve better performance (68.2 vs 63.0). We also explore an agentic task called Alfworld[1], which is an embodied multi-turn interactive task. With the same fine-tuning and RL scheme, results shown in following table demonstrate superior performance of LAMP on it:

We sincerely appreciate the reviewer's thoughtful feedback and positive recognition of the novelty of the perspective of RL from observation. We are grateful for the opportunity to address the reviewer's concerns and further clarify our methodology.

Q1: While the claims regarding improved controllability is generally supported, the main concern is that the performance gains demonstrated by LAMP are relatively marginal compared to baseline methods. Specifically, the reported improvements such as Table 1 and Table 2, appear limited in practical significance given the complexity of the proposed method.

A1: We apologize for the unclear expression. We clarify that Tables 1 and 2 show only "me too" but not the main purpose of LAMP, which is to enhance the downstream tasks using reinforcement learning.

Explanation: Table 1 shows the performance of LAMP model with behavior cloning policy, which should be called as LAMP-BC. We train LAMP-BC just to show that the LAMP structure can perform as an ordinary LLM, as presented results in Table 1. Further, Table 2 shows supervised fine-tuned performance of both the base LLM and LAMP-BC, which shows LAMP-BC can have a slightly better performance. We will revise the corresponding descriptions in Sections 5.1 and 5.2 to strengthen the presentation for the propose of these two experiments to explain this more clearly.

To provide an overview of the power of our approach, we summarize and explain our RL results below:

• For preference alignment tasks, our method achieves superior win rates and robustness via RLHF across different KL settings. In Figure 2, our LAMP model achieves an average of 64% win rate at KL=0.01 and over 90% win rate at KL=0.00. In Figure 3, our LAMP model achieves a slight performance improvement under the KL=0.00 setting compared to the standard KL=0.01 configuration, indicating a more robust alignment of our LAMP model.
• For mathematical reasoning, after fine-tuning and RL optimization, our method achieves 42.6 on MATH500, comparing with 38.2 of baseline. And achieves 68.2 on a more flexible MCTS method, comparing with 63.0 of baseline.

These results show that our method achieves significant improvements by introducing latent actions for RL training.

Q2: Additionally, the claim about the semantic interpretability of latent actions is not fully convincing. It remains unclear how these high-level semantic categories meaningfully translate into effective control or actionable improvements.

A2: We would like to clarify that we did not claim the semantic interpretability of the extracted latent actions. In fact, Figure 12 (Appendix C.2, Page 18) aims to show that latent actions have higher semantics than tokens. We will revise the claim in Section 4 to state that the connection of the latent actions to semantic is weakly correlated, avoiding cause misunderstandings.

To address your concern about the effectiveness of latent actions, we clarify that the higher-level latent actions not only enhance the reinforcement learning performance due to reduced latent action space, but also mitigate the reward hacking issue (Figure 3), which can lead to more stable reinforcement learning.

We sincerely thank the reviewer for helpful feedback and positive assessment of the novelty in our latent action control approach. We address the reviewer's concern below, with key claims highlighted in bold for clarity.

Q1: In the ablation study, why did increasing the model size and the data size lead to worse performance? This doesn't seem to obey the scaling laws of LLMs. I hope the authors would like to provide some explanation.

A1: Thank you for your question. We’d like to clarify that our findings do not conflict with the scaling laws of LLMs. The slight performance degradation observed in Figure 6 stems from the "forgetting issue" during continued pre-training of Llama-3.1-8B. This occurs because our experimental data differs from Meta’s original pre-training data, leading to partial knowledge forgetting. Similar effects have been documented in prior work [1]. To our knowledge, scaling laws [2] are primarily applied to large-scale pre-training, not continued training with limited data.

We clarify that we show Figure 6 to indicate that the original Llama-3.1-8b is a proper baseline. We will revise the descriptions to modify the explanation of this experiment in Section 5.4 to explain it more clearly. We understand that the reviewer may expect improved performance with more data. While this could be possible by enhancing LAMP with additional insights and techniques from continued training, such an investigation is beyond the scope of this paper. We appreciate your feedback and consider this an important direction for future work.

[1] Shao et al. DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv:2402.03300 (2024).

[2] Kaplan et al. Scaling Laws for Neural Language Models. arXiv:2001.08361 (2020).