摘要

尽管大型语言模型（LM）在众多任务中表现出了出色的能力，但它们经常在物理环境中无法进行的简单推理和计划，例如理解目标的持续性或计划家庭活动。这种局限源于以下事实：LM仅在书面文本中训练，而缺乏基本的实体知识和技能。在本文中，我们提出了一种新的增强LM的范式，其通过世界模型对LM进行微调，以获得多种实体知识，同时保留其通用语言能力。我们的方法在世界模型中部署了一个实体智能体，以作为物理世界（VirtualHome）的模拟器，并通过面向目标的计划和随机探索获得各种实体经验。然后，这些经验被用来微调LM去学习在物理世界中推理和行动的各种能力，例如计划和完成目标，目标持续性和追踪等。并且，其倾向于在微调期间保留LM的通用能力，以将实体知识泛化到多种任务，而不仅仅是特定模拟上。因此，我们进一步引入了经典的elastic weight consolidation（EWC）以进行选择性权重更新，并结合low-rank adapters (LoRA)，以实现较高的训练效率。广泛的实验表明，我们的方法在18个下游任务上将基础LM平均提高了64.28％。特别是，通过我们的方法增强的小型LM（1.3B和6B）匹配甚至超过大型LM的性能（例如ChatGpt）。我们的代码可在https://github.com/szxiangjn/world-model-for-language-model中找到。

1.介绍

　　语言模型（LM）在各种自然语言处理任务上表现出了令人印象深刻的表现。特别是，最近的研究表明，LM可以协助完成实体任务的决策，这证明了LM对物理世界具有一定程度的理解。但是，对于在物理环境中的许多推理和计划任务，这种理解还不够强大。如图1所示，即使是最新的大型LM，例如ChatGPT，仍然会在看似简单的查询中犯错误，例如计算指定位置上的目标数目。我们认为这是因为仅通过大规模文本语料库训练的LM没有实际的经验，例如在环境中导航，与目标互动，感知以及跟踪世界状态。因此，它们缺乏与物理环境相关的推理和计划所必需的强大而全面的实体知识。相关的研究是对LM微调，以改善特定的具体任务，从而得到特定任务的模型。
　　在本文中，我们的目标是将各种基本的实体知识和技能注入与训练的LM，同时保留模型的通用性。我们引入了一种针对LM的新训练范式，即使用来自世界模型的实体经验微调LM（E2WM）。其中，世界模型是在现实世界环境（例如VirtualHome）中模仿物理交互的实体模拟器。他们为LM提供了理解目标和环境交互并执行操作的机会。这些世界模型是我们现实世界中简化且具有成本效益的复制品，可以显着增加传统的预训练框架。我们预计，使用从世界模型中收集的实体经验上微调LM可以增强其实体知识，并保留了模型的通用性，从而增强了他们解决广泛具体任务的能力。
　　在这项工作中，我们考虑了各种实体任务所需的基本知识和技能，包括目标追踪，计划以完成给定的目标，识别其他是能体的行为等。为此，我们引入了两种方法收集来自世界的实体经验，即生成的所需的知识和技能：面向目标的计划和随机探索（图1）。具体而言，面向目标的计划旨在收集与计划和面向目标智能体的行为相关的经验，而随机探索的重点是累积涉及目标和世界状态追踪的经验。在面向目标的计划中，对于特定的活动（例如Clean Floor），模型被给定了目标（例如，IN(dust, trash can)），并因此模型应该制定一个计划来完成它。为了发现该计划，我们设计了蒙特卡洛树搜索（MCTS）来探索世界模型。然后，该过程将作为实体经验存储。在随机探索中，将一个或多个智能体部署在世界模型中以执行随机操作，而所有目标的位置和移动则是同时被跟踪的。
　　在收集了实体经验后，我们使用它们来构建一组微调任务（例如，计划生成，活动识别和追踪）。至关重要的是，为了在收集到的实体经验上微调，同时保留其原始的常识和能力，我们提出将经典的Elastic Weight Consolidation (EWC)纳入我们的训练范式。通过正则化微调损失，EWC旨在保留重要的LM参数。我们表明，EWC比流行的KL正则化更为有效。我们通过将最近的Low-Rank Adaptation (LoRA)与EWC正则化器联合起来，进一步引入了高效的低秩更新。这导致了新的EWC-LoRA更新规则，该规则大大降低了训练成本，并使我们的E2WM范式可部署在便宜的硬件上（GPU）。
　　我们使用一个虚拟家用模拟器VirtualHome来实例化世界模型，并将我们的方法应用于GPT-Neo-1.3B和GPT-J-6B模型。为了测试微调后LM的泛化能力，我们在各种未知任务上进行了评估，这些任务需要类似训练任务的实体知识来解决。此外，我们评估了模型在原始预训练数据上的性能，以确定其核心语言建模能力的保留程度。实验表明，我们的方法显着改善了已知和未知的任务的基准（例如在计划生成任务上Rouge-L从34.31提升到51.23 ，在计数任务上的准确性从30.41％到67.01％），同时不会遭受预训练数据集上的性能损失（Pile测试子集上困惑度从3.4443到3.537）。此外，用我们E2WM范式训练的小型GPT-J-6B模型，在许多任务上优于ChatGPT。实验结果表明，E2WM是使用泛化实体知识和技能微调预训练LM的机制。

2.相关工作

World Models。世界模型的概念可以解释为由机器人智能体学习的模型，以预测潜在动作的结果，或者使用物理引擎模仿现实世界的模拟器。在本文中，我们采用后者的定义。在实体AI领域，许多世界模型都是为了复制现实世界而建立的，并充当虚拟测试环境，用于在现实世界部署之前评估机器人智能体。例如，VirtualHome是由Unity3D游戏引擎实现的模拟3D家庭环境。AI2-THOR是由近3D室内场景组成，并具有更丰富的对象属性和交互类型。其他室内家庭世界模型还包括VRKitchen，CHALET，MINOS，House3D等。此外，MineCraft是一种更具挑战性和开放式的世界模型，该模型具有大量的目标和大规模的任务层次结构。在本文中，我们利用VirtualHome作为世界模型。
　　Language Model Grounding。最近的大量工作集中于将基础语言模型放置于世界模型中。其中一些通过冻结LM并利用某些提示策略或专门设计的模块。例如，Zero-Shot Planner提示LM生成活动计划并将其转化为可接受的动作。Mind’s eye提示LM使用物理引擎进行模拟，以回答物理推理问题。SayCan使用一个学习的affordance函数来协助LM选择有效的动作。DEPS提示LM，以描述，解释和生成动作计划，并与一个学习的选择器模块合并，以选择最有效的路径。还有其他工作朝着更好的下游任务性能来微调LM。 例如，, Li et al.使用具有交互决策的监督学习微调LM，而 Carta et al.通过在线强化学习来增强LM。与这些旨在目标环境的特定任务中优化LM的工作不同，我们的工作着重于通过从世界模型中获取知识来改善语言模型本身。
　　Language Model Regularization。为了促进在不失去LM语言建模能力的情况下获得新知识和技能，经常在微调过程中引入正则化。一种流行的方法是添加KL惩罚，该惩罚通过在当前训练模型和原始模型之间的输出概率上利用KL散度，以RL方式正则化LM（即通过计算策略梯度）。例如，InstructGPT使用KL惩罚来缓解奖赏模型的过度优化， Liu et al. 添加KL正则化以训练常识知识生成器。在这项工作中，我们使用弹性权重巩固（elastic weight consolidation，EWC）进行正则化。我们的实验结果表明，在保留语言建模能力和LM通用性方面应用EWC比KL惩罚更有效。

3.方法

在这项工作中，我们提出了一种新的训练范式，即使用来自世界模型（E2WM）的实体经验的微调，以将实体知识注入LM而不牺牲其通用性和语言建模能力。我们使用的世界模型是VirtualHome，这是用于家庭活动的多智能体模拟器。在VirtualHome中，一个可执行的动作步骤可以简化为$[Action]<arg>$的格式，例如$[Grab]<Apple>$。VirtualHome的世界状态由目标及其关系组成（例如， apple on table）。有关VirtualHome的详细信息可以在附录A中找到。我们首先在第3.1节中描述了如何收集世界模型中的实体经验。然后，在第3.2节中，我们通过利用收集的经验以及我们提出的方法EWC-LoRA来证明如何通过有效的知识概泛化来微调LM。

3.1 Collecting Embodied Experiences from World Model

　　在大规模人工编写的文本语料中预训练的LM在解决物理环境中的基本推理和计划方面通常会遇到困难。这是因为LM缺乏对物理世界的必要实体知识和经验理解。为了解决这个问题，我们提出利用世界模型来收集各种实体经验以增强LM。具体而言，要将不同类型的实体知识注入LM，我们介绍了两种收集经验的方法：面向目标的计划和随机探索。图2说明了两种方法。
　　Goal-oriented Planning。一项重要的实体技巧是计划和完成一个特定目标，例如正确放置餐具以布置桌子。为了赋予LM这种能力，我们提出了面向目标的计划。该方法旨在生成面向目标的经验，因此这对于在世界模型中利用技能和任务计划能力执行一系列活动非常有用。为此，我们收集了各种活动及其相应的目标。正式地，世界模型中活动的目标被定义为描述目标世界状态的一组谓词。例如，一个活动是set up table,，其目标可以是ON(fork, table);ON(plate, table)，这意味着应将叉子和板放在桌子上以实现活动。有关谓词和目标定义的更多详细信息可以在附录A中找到。如图2所示，在面向目标的计划中，我们设计了一个蒙特卡洛树搜索（MCTS）计划器，以通过动作空间搜索并找到一个计划，即一个动作序列，以实现目标。成功MCTS的关键是奖赏设计。在每个时刻，如果满足至少一个目标谓词，MCTS计划器将获得+2的奖赏，并且实现的目标谓词将被从目标中删除。这样可以确保计划器不会反复执行相同的动作来获得奖赏，而是专注于实现其余未实现的目标。此外，每次步骤之后，都会受到-0.1的惩罚，以阻止计划器在实现目标过程中采取不相关的动作。最后，我们将计划过程作为实体经验存储。
　　Random Exploration。在实际情况下，人类不仅通过完成任务来获取新知识，而且还通过随机探索周围环境，例如随机观察/跟踪对象并了解其属性来学习新知识。为了模仿这一学习过程，我们提出了另一种方法，即随机探索。通过简单地探索世界模型，出现了新的实体经验，这些经验涉及高级认知能力，包括目标持续性和跟踪，即使目标不在视野中，智能体也会观察并跟踪它们。然后，会收集经验来微调LM。具体而言，该方法在世界模型中部署一个或多个智能体，然后漫无目的并随机执行动作。如图2所示，多个智能体在同一环境中，相互交互或在同一目标上执行不同的动作，从而模拟复杂情况。在探索过程中，记录了世界模型中所有目标的移动路径和最终位置。然后将整个过程作为实体经验来捕获。

3.2 Finetuning LMs with Embodied Experiences

有多种方法可以利用收集到的实体经验进行LM微调，例如有监督学习和强化学习。在这项工作中，我们将它们用于有监督学习，以简化并提高训练效率。具体而言，面向目标的计划经验以两种格式编译为数据示例：计划生成和活动识别。如图2所示，在计划生成中，给定一些相关目标状态作为初始条件，LM必须逐步生成动作顺序以实现目标活动。在活动识别中，该模型需要识别其计划名称。从随机探索获得的经验也将转变为两个自监督的任务：计数和目标路径跟踪。这两个任务的示例可以在图2中看到。具体而言，对于计数，LM的任务是在智能体执行相关和无关紧要的动作和随机安排的目标后，在特定位置识别目标的数目和名称。在目标路径跟踪中，模型的任务是输出由不同的智能体获取的目标的移动路径。所有任务均经过交叉熵损失训练。假设$\text{x}$是输入（例如，计划生成中的初始条件），并且 y = { y 1 , . . . , y M } \textbf y=\{y_1,...,y_M\} y={y1​,...,yM​}是标签（例如，逐步动作序列），我们通过将不同的权重分配给不同的任务来微调LM：
$${\mathcal{L}}_v=\sum _{v\in V}{\alpha }_v\sum _{m=1}^{M}logP(y_m|{\text{y}}_{<m},\text{x})\text{\hspace{1em}(1)}$$
其中$\mathcal{L}$是损失函数；$V$是任务集；${\alpha }_v$是任务$v$的权重。与Flan-T5类似，$\text{x}$是一个提示，其由一个任务指令和采样的上下文演示格式化。我们在附录A中提供所有提示。
　　Efficient Finetuning with Preserved Generality。但是，简单地微调LM有两个关键问题。第一个是LM很容易地拟合到下游任务，从而导致其他任务的性能退化。这偏离了我们的目标，即模型应泛化到各种获得知识的任务。另一个问题是，对整个LM进行微调是资源密集且耗时的，尤其是当LM非常大的时候。为了克服问题并促进与世界模型的持续有效的知识获取，我们提出使用具有弹性权重巩固（EWC）的低秩适配器（LoRA），我们称为EWC-LoRA。
　　EWC是一种基于正则化的方法，通常用于持续学习领域。它计算一个Fisher矩阵，以估计每个参数对任务的重要性，然后使用它来对新任务进行正则化训练。正则化项有助于限制新任务的参数更新，以避免忘记先前的知识。令$U$表示预训练任务集，$V$是微调任务集。在[2]之后，我们有：
$$F_{i,i}=\frac{1}{N}\sum _{j=1}^N(\frac{\partial {\mathcal{L}}_U^{(j)}}{\partial {\theta }_{U,i}^{\ast }}{)}^2,\text{\hspace{1em}(2)}$$
$$\mathcal{L}(\theta )={\mathcal{L}}_V(\theta )+\lambda \sum _i{F}_{i,i}(\theta -{\theta }_{U,i}^{\ast }{)}^2,\text{\hspace{1em}(3)}$$
其中$\mathcal{L}$是损失函数，$F$是Fisher矩阵，$\lambda$是超参数，$i$和$j$分别是参数和数据样本的索引，$\theta$和${\theta }_U^{\ast }$分别是当前训练的参数和在任务$U$上训练后固定的参数。请注意，公式3中的第一项${\mathcal{L}}_V(\theta )$在公式1中进行了计算，第二项是EWC正则化器。在等式2中，通过计算在任务$U$上损失的梯度的平方和来计算Fisher矩阵，这表明了每个参数对任务$U$的重要性。在$V$训练时，通过使用EWC，LM学会了适应新任务，而不会遗忘预训练的任务，这迫使其从微调任务中理解和消化新知识，而不是对它们过拟合。
　　但是，EWC是时间和内存都不高效的。首先，它需要对整个大型LM参数组合进行微调。此外，该方法涉及创建与LM相同大小的Fisher Matrix，这导致内存开销是原始大小的三倍。 这使得应用大型LM特别具有挑战性。为了减轻问题，我们提出将EWC与低秩适配器（LoRA）相结合，这是一种参数有效的调整方法。LoRA冻结了预训练的模型权重，并将两个可训练的低秩矩阵注入模型的每一层。假设$W,W^{\ast }\in {\mathbb{R}}^{r\times d}$分别是训练好的权重矩阵和冻结的权重矩阵。$B\in {\mathbb{R}}^{r\times k}，A\in {\mathbb{R}}^{k\times d}$是两个低秩矩阵，满足$k\ll min(r,d)$。 然后，LoRA的公式可以写为$W=W^{\ast }+BA$。假设$H$是$BA$的展开。值得注意的是，我们发现等式3中的${\theta }_i$是$W$的元素，${\theta }_{U,i}^{\ast }$是$W^{\ast }$的元素。因此，${\theta }_i-{\theta }_{U,i}^{\ast }$是$H$的元素。因此，我们可以将方程3转换为EWC-LoRA方法的最终公式：
$$\mathcal{L}(\theta )=L_V(\theta )+\lambda \sum _i{F}_{i,i}{h}_i^2\text{\hspace{1em}(4)}$$
其中$h_i={\theta }_i-{\theta }_{U,i}^{\ast }$是$H$的第$i$个元素。重写等式的好处之一是，我们不再需要像原始EWC一样存储训练好的LM权重矩阵，这可以节省大量内存空间。此外，我们只需要更新B和A，这也降低了内存需求并导致训练速度更快。令人惊讶的是，如实验所示，将LoRA添加到EWC中可以进一步减轻灾难性遗忘和过拟合的问题。这与先前的结论相吻合，即限制优化问题的维度可以减轻灾难性的遗忘。