线性循环模型,例如Mamba,以及线性注意力机制,均拥有一个显著的特点:它们能够有效处理极长的序列,这一能力对于长上下文推理任务来说至关重要。
这恰恰是它们相对于 Transformer 的显著优点——Transformer 由于受到上下文窗口大小的限制,且在序列长度方面的计算复杂度为二次方,因而成为制约其性能的瓶颈。
曾经,循环模型在性能上存在较大短板:尤其是在处理较短的序列时,其表现通常不及 Transformer。但近年来,得益于架构上的多项创新,循环模型的性能得到了显著增强,在众多任务中已经能够与 Transformer 相当,甚至已经广泛应用于多个工业领域,比如音频建模和代码自动补全等。
近期的研究揭示,循环模型存在一个显著的不足之处:尽管在训练阶段内它们的表现相当出色,然而在面临超出训练序列长度的数据时,其泛化能力往往不足,导致性能显著下滑。
实际情况确实如此,以Mamba-2检查点为例,图中展示了其在不同序列位置上的性能波动(以困惑度Perplexity为指标,数值越低意味着性能越佳)。观察可知,当序列位置超出训练上下文限制后,这些模型几乎失去了效用:它们无法实现长度的泛化能力。
这引发了一个问题:目前通用的循环模型在处理长序列时效果不佳,并且与 Transformer 相比,在短序列上的效率提升并不显著;换言之,它们在两个关键方面似乎都不尽如人意。
那这是否意味着循环模型就毫无用处了呢?
当然不是!
本研究由卡内基梅隆大学和Cartesia AI的研究团队完成,他们证实了经过简单的训练调整,循环模型能够实现长度的广泛适应。仅需进行约500步的训练(这大约占预训练预算的0.1%),模型便能在长达256k的序列上实现泛化效果。
因此,循环模型并非存在根本性的缺陷,它所蕴含的潜力尚未得到充分的挖掘和发挥。
值得关注的是,Mamba项目的一名核心成员Albert Gu同样参与了该研究。在2023年,他与Karan Goel、Chris Ré、Arjun Desai以及Brandon Yang携手合作,共同创立了Cartesia公司。公司的目标在于打造一款具备持久记忆功能的即时智能系统,该系统能够在任何地点稳定运行,这与本文探讨的主题紧密相连。
为什么循环模型不能进行长度泛化?
函数f与g与位置t无关,故循环模型在理论上是能够灵活应对各种序列长度的。然而,为何在t值较大时,它们却会出现问题呢?
本文揭示了状态h_t的分布会随着时间推移而发生改变。因此,即便在某个时间点T之前,g和f运行状况良好,但在t大于T的时刻,h_t的状态可能会出现显著差异,进而使得模型难以输出准确结果。事实上,图示中Mamba-2的状态范数随着时间的推移呈现出明显的增长趋势。
这正是因为循环模型在长度泛化方面存在困难:在处理序列长度超出训练范围的情况下,模型将遭遇训练期间未曾遇到的h_t状态,故而它并未掌握应对这些新状态的方法。
鉴于这一发现,本研究构建了一个全新的理论框架——即未挖掘状态假说,旨在阐述循环模型在处理长度泛化问题时出现失误的根本性成因。
关于未被触及的领域假说(Uncharted Territory Hypothesis)。
循环模型若仅针对所有可能状态分布的一部分进行训练,那么它们在长度泛化方面将面临困难——换句话说,模型仅掌握了在有限的训练步数中可能出现的状态,却未曾接触过在无限时间状态递推后可能出现的状态分布。
随着训练时间的延长,模型可能会对这一有限的子状态集产生过度拟合,因此在面对更长的序列时,其表现不尽如人意。这是因为它们将遇到未曾经历过的状态分布,进而导致性能的下滑。
训练干预,使长度泛化
未探索状态假说提出,若要实现长度的泛化能力,无需对模型的架构或运行机制进行调整,关键在于让模型在训练过程中接触到更多样化的状态分布,特别是那些在长序列状态递推过程中自然形成的分布。
因此,最直接的手段是让模型在更长的序列中进行训练,然而在实际情况中,这通常难以实现,原因有:
所以,我们必须探索更为高效的训练途径,以涵盖这些状态分布,进而增强模型的长度泛化水平。
实现长度泛化的方法是:对初始状态进行干预
通常情况下,现代模型架构默认将初始状态设为h_(-1)=0,而本研究则对这一初始状态h_(-1)进行了四种基础调整。这四种训练时的调整手段,实际上是从四个不同的概率分布中选取h_(-1)的初始值。
图中呈现了Mamba-2模型在完成500次迭代训练(约占预训练总迭代次数的十分之一)之后,针对不同干预策略所呈现的效果。
核心发现 1:SP 与 TBTT 机制可实现长度泛化
SP与TBTT这两种干预手段,能够让模型在远远超出训练序列长度的情况下实现良好的泛化能力。从这个现象中我们可以看出,长度泛化在循环模型中通过简单的训练干预是相对容易实现的。
请注意,结果只达到了原始预训练预算的 0.02% !
核心发现二指出,我们能够根据干预措施的效果,来推测循环模型状态分布的特点。
在370M参数规模的模型中,随机噪声的干预并未达到预期的长度泛化效果,相对而言,拟合噪声却表现出了良好的效果。这一现象暗示,对于该370M模型而言,其可达状态的分布无法仅通过固定方差的高斯分布进行近似,而需要借助在每个层和每个注意力头中应用拟合方差的独立同分布高斯分布来实现近似。
尽管在1.3B模型中未能成功实现长度的泛化,这表明大模型的状态可能在其各个元素间存在着更为复杂的相互依赖,因而无法再采用简单的独立同分布(IID)模型进行描述。
除此之外,这些干预措施还有助于解决之前所呈现的状态范数随时间逐渐增大的问题,确保模型在各个时间点输出的状态范数保持一致,进而增强了模型的整体稳定性。
长上下文任务的表现
本研究发现,这些干预手段确保了在训练时文本长度T增加后,模型性能依然稳定,然而,目前尚不明确这些手段是否具备在处理涉及超过T个位置token间关系任务时的泛化能力。
人们或许会提出疑问,这些干预手段是否仅仅是通过限制模型在超出训练时上下文长度的情况下进行推理,来达到增强其长度鲁棒性的目的?
这种机制与滑动窗口注意力相似,模型无法对超出滑动窗口范围的token进行推理;在所有评估中,当t大于T时,该模型将维持稳定的性能表现;然而,它无法应对那些需要基于长上下文进行推理的任务。
在本文的研究中,我们通过在三个不同的长上下文任务中进行的实验,证实了这些干预手段确实能够有效实现长度的泛化效果。
BABILong
BABILong是一项颇具难度的基准测试,旨在评估模型对常识知识的掌握程度以及其识别文本中远距离依赖关系的能力。
通过观察下方的图像,我们可以发现,在样本量较少以及微调的条件下,状态传递(State Passing)这一机制显著提升了模型在序列长度为2048时的长度泛化能力。
因此,状态传递不仅对于缓解已构建语言模型中出现的困惑度扩散问题具有积极作用,而且还能显著提升其处理长上下文推理任务的能力。
密码检索任务
模型在处理密码检索任务时,需在较长的文本内容中,于指定深度的位置查找一个由五个数字组成的密码。
下图中呈现了Mamba-2 370M和780M在三种不同配置中的运行情况:包括零样本学习、常规微调以及通过引入拟合噪声来进行的微调。在经过拟合噪声微调后,模型得以有效运用超过2048个位置(即训练上下文长度)的token关系。尤其是780M模型,它能够出色地完成长度达到256k的序列中的密码搜索任务。
合成复制任务
合成复制任务要求模型复制一个任意的 token 序列。
该表揭示了在训练阶段引入状态传递技术,显著增强了模型在处理长度超过三倍序列时的验证效果。由此可知,状态传递有助于模型达成对较长序列的泛化能力,并成功应对了训练阶段未曾遭遇的、更为复杂的长上下文问题。
深入探讨循环模型如何处理上下文
本文阐述了干预初始状态可确保模型在长度方面具有稳定性,同时提升了模型处理长文本序列的能力。基于这些研究成果,本文进一步提出了一种评估方法,旨在深入揭示序列模型在上下文处理方面的机制。
在理想状态中,对于文本建模而言,我们期望模型能集中精力于贴近的语境,而非过分沉迷于那些相隔甚远的标记。那么,我们该如何对此行为进行量化评估呢?
尽管每个标记都会对模型的结果造成一定影响,然而从直观感受来看,我们普遍认为紧邻的标记理应拥有更为显著的影响力。
然而,观察到 EffRem 曲线在迅速上升后趋于平稳。这一现象显然表明存在问题:在 T=8192 的时刻,对于下一个 token 的预测结果,不应当因为模型是仅基于最近的 token 还是整个序列进行判断,而产生显著的波动。
在自然语言处理中,模型应当主要参考最近的语境信息,而早期词汇不应导致预测结果发生剧烈变化,尤其是当这种变化使得两个输出概率分布的差异接近于1时。这表明,模型在序列起始部分的token上受到了过度的关注。
状态传递修正了有效记忆
经过状态传递的后续训练,EffRem 曲线呈现出上升趋势,这说明模型对远程 token 的权重值较低,并且这一权重值正逐渐提升。特别是,紧邻的上下文 token(比如句子中的前一词)对后续 token 的预测起到了关键作用,这恰好符合文本建模中应有的行为特征。
总的来说,凭借高效的记忆方式,我们能够确保模型在处理状态传递时,优先考虑最近的上下文信息,从而避免被较远历史中的token无谓地影响。