[论文解读] Amortizing intractable inference in large language models

J. Edward Hu, Moksh Jain+5 more|arXiv (Cornell University)|Oct 6, 2023

Topic Modeling被引用 4

一句话总结

本文提出通过使用生成流网络（GFlowNets）微调大型语言模型（LLMs）实现近似贝叶斯推断，以实现对难以处理的后验分布的高效、多样化采样——例如链式思维推理和约束生成任务中所需的分布。该方法在数据效率和样本多样性方面优于监督微调（SFT）和PPO，在使用50个种子推理过程的情况下，于分布内整数运算任务上实现了95.2%的准确率。

ABSTRACT

Autoregressive large language models (LLMs) compress knowledge from their training data through next-token conditional distributions. This limits tractable querying of this knowledge to start-to-end autoregressive sampling. However, many tasks of interest -- including sequence continuation, infilling, and other forms of constrained generation -- involve sampling from intractable posterior distributions. We address this limitation by using amortized Bayesian inference to sample from these intractable posteriors. Such amortization is algorithmically achieved by fine-tuning LLMs via diversity-seeking reinforcement learning algorithms: generative flow networks (GFlowNets). We empirically demonstrate that this distribution-matching paradigm of LLM fine-tuning can serve as an effective alternative to maximum-likelihood training and reward-maximizing policy optimization. As an important application, we interpret chain-of-thought reasoning as a latent variable modeling problem and demonstrate that our approach enables data-efficient adaptation of LLMs to tasks that require multi-step rationalization and tool use.

研究动机与目标

• 为解决自回归LLM在可 tractably 查询自回归采样之外的知识方面的局限性。
• 在序列续写、填空和约束生成等任务中，实现对难以处理的后验分布的高效、多样化采样。
• 与标准微调和强化学习相比，提升少样本推理和工具使用适应中的数据效率和样本多样性。
• 将链式思维推理建模为可应用近似贝叶斯推断的隐变量建模问题。

提出的方法

• 使用GFlowNets微调预训练LLM，GFlowNets是一种寻求多样性的强化学习算法，可训练策略以按与奖励函数成比例的概率采样对象。
• 在训练期间使用LLM自身的似然作为奖励函数来评估序列，从而实现自监督奖励建模。
• 以预训练LLM初始化GFlowNet策略，并通过匹配输入、推理过程和输出的联合概率的奖励目标继续训练。
• 使用种子推理过程的缓冲区以改善探索和初始策略质量，尤其在低数据场景下。
• 采用温度调度和课程学习以稳定训练并改善收敛性。
• 使用LoRA进行高效微调，并调整超参数以优化样本多样性和奖励最大化。

Figure 1: Left: Three problems of reasoning in language – sentence infilling, chain-of-thought reasoning, and problem-solving with external tool use – can all be seen as instances of the latent variable model at the top left, where an input ( $X$ ) generates the output ( $Y$ ) via a latent variable

实验结果

研究问题

• RQ1 通过GFlowNet微调实现的近似推断是否能有效且多样化地从LLM中的难以处理后验分布中采样？
• RQ2 在推理任务中，GFlowNet微调相较于监督微调和PPO在数据效率和样本多样性方面表现如何？
• RQ3 在监督有限的情况下，GFlowNet微调是否能有效学习生成有效且正确的算术推理步骤？
• RQ4 在低数据少样本推理场景中，使用种子推理过程缓冲区是否能显著提升性能？
• RQ5 GFlowNet微调能否减少PPO-based策略学习中常见的模式崩溃和过度优化问题？

主要发现

• 当使用50个种子推理过程时，GFlowNet微调模型在分布内整数运算任务上实现了95.2%的测试准确率，显著优于零样本提示和基线微调方法。
• 在使用50个种子推理过程的情况下，GFlowNet模型在分布外（OOD）算术问题上达到了75.4%的准确率，表明其在训练分布之外具有强大的泛化能力。
• 该方法生成了有效、逐步的推理过程，能够正确评估表达式，而PPO则常生成重复或无效序列。
• 随着种子推理过程数量的增加，性能持续提升，当不使用任何种子示例时，准确率仅为22.6%，凸显了初始监督对探索的重要性。
• PPO模型生成了高奖励但无效的推理过程（例如重复表达式），而GFlowNets则生成了正确且多样的推理链，表明其模式崩溃现象更少。
• 消融实验确认，种子推理过程对性能至关重要，因为若无种子，性能会急剧下降，表明在缺乏先验知识的情况下探索极具挑战性。

(a) Base model 50.5% of samples are valid numbers.