前言TTT出来有一段时间了让我确定要写TTT解读的是源于我司LLM论文100篇课程群里的一学员辰子说“校长 最近的TTT考不考虑讲一下”以及微博上也有朋友问“老师TTT 怎么说”故当时想着解读完mamba2之后则解读open-television、我司7方面review微调gemma2再接下来是TTT、nature审稿微调、序列并行、Flash Attention3、vLLM(Efficient memory management for large language model serving with pagedattention)..如今虽然mamba2的解读还没完全修订完但“open-television、我司7方面review微调gemma2”都解读的差不多了故今天开写TTT对应论文为Learning to (Learn at Test Time): RNNs with Expressive Hidden States比我科研水平高的人大有人在但比我写的大模型技术博客还要更通俗易懂的则寥寥无几究其原因还是在于比他人花更多的时间和心思对于他人写一篇博客可能只愿意花1-2天时我愿意花1-2周对于他人改一篇博客可能只愿意花1-2周时我愿意花1-2个月核心动力还是来源于写博十多年下来积累的习惯尽我最大努力让最广大的读者可以最大程度的最快理解第一部分 TTT的关键探索与关键方法1.1 TTT提出之前的背景与关键探索1.1.1 困境RNN的线性复杂度与面对长下文时的局限虽说相对于transformer的二次复杂度RNN有着线性的复杂度当然如果序列长度比较短那么无论是二次还是线性其实区别不大只有在序列长度比较长即长上下文时RNN这个线性的复杂度才有比较高的价值或意义那到底多长之后会使得RNN这个线性复杂度具有极高的价值或意义呢如下图所示这个上下文长度是在8k之后可问题是一旦上下文足够长后现有的RNN如Mamba在实际利用这些额外信息时会遇到困难(On the other hand,once context is long enough, existing RNNs such as Mamba struggle to actually take advantage ofthe extra information being conditioned on)为何RNN在面对长下文时处理起来会比较困难呢原因在于与自注意力不同RNN 层必须将上下文压缩到一个固定大小的隐藏状态中。作为一种压缩启发式方法其更新规则需要从成千上万甚至可能数百万个 token 中挖掘出潜在的结构和关系如下图所示先根据输入和前一时刻的隐藏状态计算出最新的隐藏状态在此之后便可以根据最新的隐藏状态预测出了至于RNN的详细介绍详见此文如何从RNN起步一步一步通俗理解LSTM1.1.2 对比RNN与Transformer的各自优劣其实所有序列建模层都可以从将历史上下文存储到隐藏状态的角度来看比如RNN层——如LSTM [33]、RWKV [56]和Mamba [26]层——在时间上将上下文压缩到一个固定大小的状态中如下图所示(如你所见一个通用的序列建模层表示为一个根据更新规则转换的隐藏状态。所有序列建模层都可以看作是该图中三个组件的不同实例初始状态、更新规则和输出规则)简单的RNN和TTT层都将增长的上下文压缩成固定大小的隐藏状态(Both the naive RNN and TTT layer compress the growing context into a hidden state of fixed size, therefore their cost per token stays constant)然这种压缩有好有弊好的是其输入token到输出token的映射十分高效即每个token的更新与输出的时间都是常量时间成本可控mapping an input token xt to output token zt is efficient, because both the update rule and output rule take constant time per token不好的是RNN层在长上下文中的性能受到其隐藏状态的表达能力的限制——即表达能力不强(the performance of RNN layers in long context is limited by the expressive power of its hidden states st)而只有隐藏状态表达的好才能让最终的整体性能好否则 则难说了..自注意力机制的隐藏状态(通常称为键值KV缓存)随着上下文增长因此每个token的成本也在增长(随着线性增长)我是中国我q1k1q2k1q3k1q4k1爱q1k2q2k2q3k2q4k2中q1k3q2k3q3k3q4k3国q1k4q2k4q3k4q4k4总之这里的隐藏状态明确的存储了所有历史上下文而不进行压缩好的是表达能力强不好的是成本增加快所以现在的问题变成了我们既希望将成千上万甚至数百万的token压缩到一个隐藏状态中且该隐藏状态还能有效的捕捉它们的底层结构和关系(we need to compress thousands or potentially millions of tokens into a hidden state that can effectively capture their underlying structures and relationship)从而兼顾高效与质量1.2 TTT更新隐藏状态的关键输出规则和更新规则过去的一年半LLM火爆全球其通过自监督的下一个词预测任务——即下一个token预测进行训练它们的权重可以被视为互联网上现有知识的压缩存储形式(所以对于咱们大模型开发者而言预训练或微调后得到的模型权重是最关键的产物如果一个模型的权重没开源则不算开源)通过查询 LLMs可以从它们的权重中提取知识更重要的是LLMs 通常表现出对现有知识之间语义连接的深刻理解以表达新的推理片段总之种种迹象表明可以使用自监督学习来压缩历史上下文到一个隐藏状态通过将上下文变成一个无标签的数据集并将隐藏状态变成一个模型(Our key idea is to use self-supervised learning to compress the historic context x1, . . . , xt into a hidden state st , by making the context an unlabeled dataset andthe state a model)类似的可以将隐藏状态现在等同于——模型的权重这可以是一个线性模型或一个小的神经网络(the hidden state st is now equivalent to Wt , the weights of a model f , which can be a linear model, a small neural network, or anything else)而TTT通过建模为模型的权重来更新隐藏状态的过程中涉及到输出规则和更新规则对于前者输出规则很简单即为公式1直白讲输出token只是对的预测怎么预测呢基于由使用更新后的权重进行预测(the output token is just the prediction on xt , made by f with the updated weights Wt)为方便大家更好的理解我举个例子如下表格所示预测“我 是 中 国”之后的下一个token人我我是我是中 对{我 是 中 国}的压缩我是中国其中 {或是“人”、或是“的”}我是中国基于可得的预测 人对于后者更新规则是对某些自监督损失进行梯度下降的一步(记为公式2)其中学习率为从压缩的角度来看每个启发式方法都需要决定记住或忘记哪个输入而此处的记住产生大梯度的输入——直观地说就是那些让学到很多的输入此外其中一种选择loss的方法是重建本身为了重建可以将处理成一个corrupted的输入然后如下优化(记为公式3)这个表达式类似去噪自编码器 [75]需要发现维度之间的相关性以便从部分信息中重建它(f needs to discover the correlations between dimensions of xt in order to reconstruct it from partial information ˜xt)如下图所示梯度下降能够减少损失但不能将其减少到零如上图所示自监督TTT损失 ℓ在所有测试序列上的平均值形式为 x1, . . . , xT其中 T 2048对于具有125M参数的网络中的前三个TTT层一次梯度下降能够将TTT损失从 ℓ(Wt−1; xt) 降低到 ℓ(Wt ; xt)且随着 t沿测试序列进一步移动 ℓ(Wt ; xt) 也进一步从 ℓ(W0; xt) 改善另为了视觉清晰损失值已在10个时间步的滑动窗口上取平均与其他RNN层和自注意力机制一样算法将输入序列 x1, . . . , xT映射到输出序列 z1, . . . , zT可以通过使用上述隐藏状态、更新规则和输出规则编程到序列建模层的前向传递中即使在测试时新层仍然训练一组不同的权重序列W1, . . . , WT——对于每个输入序列因此称之为测试时训练(TTT)层且TTT层的前向传递也有相应的反向传递前向传递仅由标准的可微分算子组成除了梯度算子∇。然而∇只是映射从一个函数到另一函数比如到并且 ∇ℓ也是由可微算子组成的从概念上讲对 ∇ℓ 调用反向传播意味着对梯度进行梯度计算另将训练更大的网络称为外循环而在每个 TTT 层内训练称为内循环两个嵌套学习问题之间的一个重要区别是内循环梯度是相对于(即的参数)计算的而外循环梯度是相对于网络其余部分的参数计算的将其表示为(外循环参数始终用表示并带有各种下标)最后你将在下文依次看到更多的公式(公式4-7还不理解 没事继续看下文即可)输出规则(外循环)公式1公式3公式4公式5更新规则(内循环)公式2公式6公式71.2.1 对于输出规则——使用TTT层训练隐藏层网络重建x中设计损失函数考虑到TTT的最终目标是使公式1——在语言建模上表现良好故可以基于人类先验知识的自监督任务采用一种更端到端的方法——直接优化自监督任务以实现下一个词预测的最终目标具体来说在外循环中学习自监督任务从公式3中的简单重建任务开始添加了一些外循环参数使这个任务可以学习为产生从的损坏一种设计是使其成为低秩投影其中是一个可学习的矩阵总之被称为训练view由于并不是中的所有信息都值得记住因此重建标签可以是另一个低秩投影而不是这里被称为标签view其中也是可学习的因此新的自监督损失由变成如下所示(公式4)注意因为在内循环中只有被优化所以写作为loss的一个参数而是这个损失函数的超参数在外循环中与一起被优化而只是一个隐藏状态而非一个参数如下代码所示其中和作为TTT层的参数实现类似于自注意力的Key和Value参数(TTT_Layer可以像其他序列建模层一样被插入到更大的网络中对于下图左边训练网络将优化TTT_Layer中Task的参数因为两者都是nn.Modul的子类对于下图右边由于Learner不是nn.Module的子类state.model在每次调用state.train的内循环中手动更新。为了简化有时将model重载为model.parameter最后训练view的维度比少因此不能再使用方程中的输出规则那怎么办呢最简单的解决方案是创建一个测试view并将输出规则更改为(公式5)弄这么多不同view的好处是什么呢首先训练view和标签view指定了在中被压缩到并随时间向前传播的信息其次测试视图指定了可能不同的信息这些信息被映射到当前输出token并通过网络层向前传播因此为自监督任务增加了更多的灵活性1.2.2 对于更新规则——使用小批量TTT进行并行化小批量梯度下降目前的TTT有个比较明显的问题是其更新规则——公式2不能被并行化因为在两个地方依赖于一个地方是在减号前一个地方在内部而由于后者内部包含了大部分计算所以我们要重点对后者——内部做并行化首先梯度下降GD有许多变体GD的一般更新规则可以表示为(记为公式6)其中是下降方向这个公式的价值和意义在于一旦计算出便可以基于上述公式的后半部分且通过cumsum获得所有的而对于我们所想要的更新规则是在线梯度下降使用为了并行化可以对所有这些变量相对于进行计算在这种变体情况下其中被称为批量梯度下降因为与相对于的梯度在作为一个批次时是相同的不过在批量梯度下降中实际上只比多一步梯度步长「这与在线梯度下降形成对比因为在在线梯度下降距离有步之遥」因此批量梯度下降的有效搜索空间会相对比较小怎么办呢好在可以在批量梯度下降与在线梯度下降之中取个折中小批量梯度下降即将批量大小设置为相对较小的b具体如下图所示然后使用其中是前一个小批量的最后一个时间步(对于第一个小批量为0有点类似的意思)因为可以一次并行化个梯度计算进一步这个其实控制着速度与质量之间的权衡如下图所示(在TTT的所有实验中作者们选择)总之有两个潜在的渠道可以将信息从传播到其中对梯度算子做累加和(cumsum and the gradient operator)累加和始终处于活动状态但梯度通道仅在来自先前的小批量时才处于活动状态。 不同的梯度下降变体仅影响梯度通道即下降方向特别是相对于梯度的计算然而下降步骤总是从开始由于更新规则的自回归性质这与的选择无关然上面面介绍的并行化是必要的但对于wall-clock time的效率而言还不够(说白了就是速度还不够快)故接下来咱们来探讨下对偶形式现代加速器专门用于矩阵-矩阵乘法称为 matmuls。例如NVIDIA A100 GPU 包含高度优化的单元称为 TensorCores它们只能执行一种操作——将两个大小为 16 ×16 的矩阵相乘。 如果没有足够的矩阵乘法TensorCores 就会闲置A100 的大部分潜力将无法实现不幸的是即使使用小批量开发的 TTT 层仍然有很少的矩阵乘法考虑 ℓ的最简单情况其中仅针对第一个大小为 b的 TTT 小批量此外考虑作为线性模型复制公式3在时间的损失为为帮助大家更好的理解回顾一下之前的公式1、公式2、公式3以及公式6输出规则(外循环)公式1公式3公式4公式5更新规则(内循环)公式2公式6如上一节所讨论的可以并行化以下计算对于 t 1, . . . , b然而不能通过一个单一的matmul计算所有的b个相反需要 b个外积来逐个计算它们。更糟糕的是对于每个是 d × d这会比在大时产生更大的内存占用和I/O成本为了解决这两个问题他们做了一个简单的观察实际上不需要具体化只要我们能在小批量结束时计算和输出token现在用上面简化的TTT-Linear案例来演示这些计算记 X [x1 , . . . . , xb], 然后所以 Wb可以通过 matmul方便地计算。 为了计算 Z [z1 , . . . . . . . . . . . . . ., zb]我们知道(记为公式7)记和矩阵可以推导出其中掩码是具有零值的下三角掩码(类似于注意力掩码但用零代替无穷大)并且项可以从的计算中重用。 现在 ∆也可以方便地用矩阵乘法计算。将 ∆代入公式7我们得到以上称这个过程为对偶形式与之前的原始形式相对比其中 G和 W是显式物化的如前所述这两种形式在输出上是等价的第二部分 两种TTT层的变体——TTT-Linear和TTT-MLP2.1 证明具有线性模型和批量GD的TTT层等价于线性注意力2.1.1 定理1及其证明回顾一下在1.2节的开头提到 f可以是线性模型或神经网络在1.2.2小节中还讨论了更新规则的三种变体在线GD、批量GD和小批量GD这三种2 ×3组合中的每一种都会引发TTT层的不同实例化如下图所示如上图所示参数化学习器需要定义两个属性模型和优化器(左每个学习器唯一地引发一个TTT层(右)本文则提出了两种引发的TTT层TTT-Linear和TTT-MLP具有线性模型和批量GD的TTT层等价于线性注意力[41]其实在这些引发的实例化中具有线性模型和批量GD的TTT层其实等价于线性注意力[Transformers are rnns: Fast autoregressive transformers with linear attention]——一种广为人知的RNN层简而言之线性注意力 [41] 只是没有softmax的自注意力。 回顾自注意力的定义没有 softmax , 其便变成了这是线性注意力的最简单形式与其他RNN层类似它可以写成递归形式其中是隐藏状态而由于可以通过 cumsum在每个计算因此线性注意力相对于 T也具有线性复杂度定理1考虑TTT层其中作为内循环模型批量梯度下降作为更新规则并且然后给定相同的输入序列公式5中定义的输出规则产生相同的输出序列作为线性注意力且为方便大家理解特此再列一下上文介绍过的公式4、公式5、公式6输出规则(外循环)公式1公式3公式4公式5更新规则(内循环)公式2公式6其中该定理的证明如下根据公式4中的定义有且根据公式6中的批量GD定义可知将代入公式5中的输出规则可得到输出token这是线性注意力的定义2.1.2 定理2定理2考虑使用Nadaraya-Watson估计器[7, 12]定义的TTT层其中是第1.2.1节讨论的标签view从公式3中的简单重建任务开始添加了一些外循环参数使这个任务可以学习为产生从的损坏一种设计是使其成为低秩投影其中是一个可学习的矩阵总之被称为训练view由于并不是中的所有信息都值得记住因此重建标签可以是另一个低秩投影而不是这里被称为标签view其中也是可学习的并且是一个带有带宽超参数的核函数和然后给定相同的输入序列, 公式5中定义的输出规则产生与自注意力相同的输出序列对于上面的TTT层隐藏状态是或类似的处理过的训练数据列表更新规则将添加到列表中输出规则使用扫描列表在前面的章节中隐藏状态被定义为更新规则是一个梯度步骤输出规则是调用为了统一这两种构造可定义一种新的抽象称为学习者它唯一地引发了TTT层类似于标准机器学习包中的定义 [54]所有学习者都需要实现两个方法训练和预测。 现在将induced的TTT层的隐藏状态重新定义为学习者的内部存储并将更新和输出规则重新定义为训练和预测方法在这种新的TTT层定义下定理1中的参数学习器和定理2中的非参数学习器都可以包括在内比如下图便总结了在所有序列建模层的更广泛范围内TTT层的这一通用定义这种通用定义对参数学习器有一个额外的好处在参数学习器的内部存储中除了之外还可以有更多的对象例如优化器状态这也将包含在诱导的TTT层的隐藏状态中。 这种扩展允许TTT层在未来的工作中使用更复杂的优化器例如Adam [42]2.2 实现的更多细节以及骨干架构2.2.1的实例化提出了两种TTT层的变体——TTT-Linear和TTT-MLP它们仅在的实例化上有所不同对于TTT-Linear其中W是方阵对于TTT-MLP有两层类似于Transformer中的MLP具体来说隐藏维度是输入维度的4×然后是GELU激活[31]且为了在TTT期间获得更好的稳定性总是包含层归一化(LN)和残差连接。即其中可以是或2.2.2 可学习的TTT初始化在所有序列之间共享尽管后续权重对于每个输入序列是不同的。可以将作为外循环的一部分来学习而不是将其设置为0由于外循环参数总是用θ而不是W表示将别名分配给它。 在实践中与重建视图相比增加的参数量可以忽略不计因为它的输入和输出都是低维的。根据经验观察到学习显著提高了训练的稳定性2.2.3 可学习的先回顾下1.2.2 对于更新规则——使用小批量TTT进行并行化小批量梯度下降目前的TTT有个比较明显的问题是其更新规则——公式2不能被并行化因为在两个地方依赖于一个地方是在减号前一个地方在内部而由于后者内部包含了大部分计算所以我们要重点对后者——内部做并行化首先梯度下降GD有许多变体GD的一般更新规则可以表示为(记为公式6)其中是下降方向这个公式的价值和意义在于一旦计算出便可以基于上述公式的后半部分且通过cumsum获得所有的而对于我们所想要的更新规则是在线梯度下降使用为了并行化可以对所有这些变量相对于进行计算在这种变体情况下其中被称为批量梯度下降学习率通常是梯度下降中最重要的超参数因此尝试在外循环中学习内循环学习率如之前介绍过的公式6为了增加灵活性可以使成为输入token的函数(因此在时间上有所不同)具体来说设计其中可学习向量是一个外循环参数是sigmoid函数标量是基础学习率对于TTT-Linear设为1对于TTT-MLP设为0.1。当然也可以解释为的一个门2.2.4 骨干架构默认使用Mamba骨干网络将任何RNN层集成到更大的架构中的最干净方法是直接替换Transformer中的自注意力机制在这种情况下称为骨干然而现有的RNN如Mamba [26] 和Gri n [18] 都使用与Transformer不同的骨干。最显著的是它们的骨干在RNN层之前包含时间卷积这可能有助于收集跨时间的局部信息且在实验了Mamba骨干后发现它也提高了TTT层的困惑度因此TTT将其纳入TTT的方法中详见下图如上图所示左一个残差块Transformer的基本构建块。序列建模块被实例化为两种变体Transformer骨干和Mamba骨干中Transformer骨干中的TTT层。 在 O之前的LN来自NormFormer [Normformer: Improved transformer pretraining with extra normalization]右受Mamba [26]和Griffin[Griffin:Mixing gated linear recurrences with local attention for efficient language model]启发的骨干中的TTT层。 遵循这两种架构这里是GELU [31]为了在不改变嵌入维度的情况下容纳门控的额外参数可简单地将和组合成一个投影实际实现时由于Transformer和Mamba使用不同的骨干网络而TTT-Linear和TTT-MLP总是默认使用Mamba骨干网络除非另有说明第三部分 实验对TTT效果的评估接下来通过与两个基线——Transformer和现代RNN Mamba进行比较来评估TTT-Linear和TTT-MLP且主要代码库基于EasyLM [25]这是一个用于在JAX中训练和服务LLM的开源项目。另所有实验都可以使用这里《https://github.com/test-time-training/ttt-lm-pytorch》提供的公开代码和数据集进行复现在数据集的选择上根据Mamba论文 [26]在Pile [24] 上进行标准实验使用2k和8k的上下文长度这是一个用于训练开源LLM的流行文档数据集 [9]然而Pile包含的长度超过8k的序列很少 [19]。 为了评估长上下文的能力我们还在1k到32k范围内以2×递增的上下文长度进行实验使用一个名为Books3的Pile子集该子集已被广泛用于训练长上下文的LLM [49, 3]且为了确保评估公平作者们在可能的情况下严格遵循Mamba论文中的评估协议对于每个评估设置例如数据集、上下文长度和方法我们实验了四种模型大小125M、350M、760M和1.3B参数。 对于Mamba相应的大小是130M、370M、790M和1.4B因为Mamba不遵循Transformer配置所有模型都使用Mamba论文中描述并在附录C中重现的Chinchilla配方9进行训练。他们的Transformer基线基于Llama架构[73]也遵循Mamba论文中的基线值得一提的是与Mamba论文中的唯一不同之处在于分词器。 Mamba论文在各种实验中使用了两种不同的分词器——GPT-2和GPT-NeoX。 为了保持一致性TTT坚持使用单一分词器并选择了Llama 2的分词器没有对混合架构例如 Grin [18]进行实验因为基线不是混合架构。 虽然使用自注意力和TTT层的混合架构可能会提高性能但它们会降低学术评估的清晰度3.1 短上下文the Pile3.1.1 Transformer与mamba、TTT-linear、TTT-MLP的相互PK如下图所示「当一个图同时包含Transformer骨干和Mamba骨干时则分别用(T)和 (M)表示」在2k上下文中TTT-Linear (M)、Mamba和Transformer具有可比的性能因为这些线条大多重叠。 在大FLOP预算下TTT-MLP (M)的表现略差。 尽管TTT-MLP在每个模型大小上都比TTT-Linear具有更好的困惑度但在FLOPs上的额外成本抵消了这一优势在8k上下文中TTT-Linear (M)和TTT-MLP (M)的表现显著优于Mamba这与在2k时的观察结果相反。 即使是具有Transformer骨干的TTT-MLP (T)在大约1.3B时表现也略好于Mamba且在本文中观察到的一个稳健现象是随着上下文长度的增加TTT层相对于Mamba的优势变得更大在8k上下文中Transformer在每个模型规模上仍然具有良好的如果不是最好的困惑度但由于FLOPs的成本其曲线不具竞争力总之TTT-Linear在2k上下文中表现与Mamba相当而在8k上下文中表现更好3.1.2 骨干网络的影响具备mamba骨干网络的TTT-Linear可以比肩TTT-MLP将TTT层从Mamba骨干网络切换到Transformer骨干网络有两个影响首先迄今为止具有Mamba骨干网络的TTT层在评估中表现更好Switching the TTT layers from Mamba backbone into Transformer backbonehas two effects. First,TTT layers with Mamba backboneperform better in our evaluations so far.其次使用Mamba骨干网络时TTT-MLP(M)最多只能与TTT-Linear(M)相当但使用Transformer骨干网络时TTT-MLP(T)明显更好Second, with Mamba backbone, TTT-MLP at best is only comparable to TTT-Linear; but with Transformer backbone, TTT-MLP is clearly better.我们假设当序列建模层具有较不具表现力的隐藏状态时Mamba骨干中的时间卷积会更有帮助(We hypothesize that the temporal convolutions in the Mamba backbone help more when the sequence modeling layer has a less expressive hidden state.)所以虽说线性模型的表现力不如MLP但线性模型可以从mamba的卷积中受益更多才使得具备mamba骨干网络的TTT-Linear可以比肩TTT-MLP(The linear model is less expressive than the MLP, therefore benefits more from the convolutions.We will revisit this hypothesis in the next subsection)即线性 mamba ≈ MLP即mamba的卷积部分可以弥补线性模型的劣势3.2 长上下文Books如下图所示在Books的2k上下文中Pile 2k的所有观察结果仍然成立除了Mamba现在表现略优于TTT-Linear(而在Pile 2k中它们的线大致重叠在32k上下文中TTT-Linear (M)和TTT-MLP (M)的表现都优于Mamba类似于Pile 8k的观察结果。 即使使用Transformer骨干的TTT-MLP (T)在32k上下文中表现也略好于Mamba在1.3B规模下TTT-MLP(T)仅略逊于TTT-MLP (M)。 如前所述由于缺乏清晰的线性拟合很难推导出经验缩放定律。然而TTT-MLP (T)的强劲趋势表明Transformer骨干可能更适合于评估之外的更大模型和更长上下文第四部分 与相关工作的对比总结4.1 TTT-MLP优于TTT-Linear而TTT-Linear又优于MambaMamba是众多结构化状态空间模型之一 [27, 21, 57, 18]。 这些模型中的隐藏状态是一个向量类似于LSTM在TTT-Linear或TTT-MLP中隐藏状态是一个矩阵或两个矩阵因此更大如下图所示我们发现TTT层可以利用其更大的隐藏状态在长上下文中压缩更多信息其中TTT-MLP优于TTT-Linear而TTT-Linear又优于Mamba与TTT-Linear类似RWKV [55, 56]、xLSTM [5]和门控线性注意力(GLA)[79]也具有矩阵隐藏状态这些状态继承自线性注意力 [41]。 现代RNN如GLA使用块状并行来提高硬件效率因此块内的标记可以通过matmul而不是cumsum来处理然而块状并行并没有改变模型的表达能力因为所有时间依赖性仍然等同于 cumsum相比之下小批量TTT允许跨小批量的更复杂的时间依赖性。每个隐藏状态依赖于其小批量内的先前仍然通过 cumsum但也通过梯度算子依赖于先前小批量中的如之前这图所示小批量TTT在表达能力和硬件效率之间实现了权衡因为较小的批量大小b会以更高的延迟为代价带来更好的困惑度。 这种权衡是TTT的一个独特且重要的特性如下表所示中间批量大小b 16显著优于 bT完全cumsum如上表所示这里所有模型都有 125M 参数并按照小节 3.1 中的配方进行训练最后一行困惑度为 11.09是下图中 TTT-Linear 的最终结果从小节 2.1 中讨论的等效性开始可学习的略有损害但下面的行在没有它的情况下无法稳定训练最大的改进来自小批量 TTT(从 b T 2048 改为 b 16)第二个来自于实例化内部模型与LN和残差连接如果没有TTT的概念框架这两种设计都将很难遇到4.2 测试时学习——Learning at Test Time在机器学习中测试时学习的概念有着悠久的历史。 这种概念最早的版本之一被称为局部学习Bottou 和 Vapnik [10]对于每个测试输入在做出预测之前先对其邻居进行训练。 这种程序已被有效地应用于从SVM [81]到现代LLM [29]的各种模型测试时学习的另一个早期版本被称为传导学习 [22]。 Vladimir Vapnik [74]提出的传导原则是“...得到你真正需要的答案而不是更一般的答案。”传导学习的实际应用使用测试数据来为SVM的边界添加约束 [39, 17]。 然而传导学习通常需要多个测试实例才能在经验上有效不像许多测试时训练的实例化只需要一次一个测试实例(图像、视频或自然语言序列)在计算机视觉中测试时学习的理念已经应用于面部检测 [38]、物体检测 [53]、图像超分辨率 [65]和 3D 重建 [50] 等应用领域数十年最近同样的理念也被应用于自然语言处理领域在那里它被称为动态评估 [44, 45]。 基本方法是直接在测试序列上微调语言模型这通常以prompt的形式出现接下来我们详细讨论两个相关的工作方向测试时训练和快速权重4.2.1 测试时训练——Test-Time Training测试时训练 (TTT) 的核心思想是每个测试实例定义其自身的学习问题其中该测试实例本身是泛化的目标 [69]具体来说对于每个测试实例x传统做法是使用预测器预测该预测器是针对所有训练实例的平均优化的TTT首先通过定义一个学习问题然后在上训练一个模型(通常以f作为初始化)并预测由于测试实例没有标签学习问题只能通过自监督任务来制定。 先前的工作表明使用重建的TTT显著提高了性能尤其是在异常值上[23]当在以流的形式到达的视频帧上进行测试并且TTT是自回归的时改进变得更加显著[76]因为是在过去的帧上训练的。 自回归连接使得[76]与TTT的论文最相关从概念上讲TTT与先前工作的最大区别在于TTT的重建任务是在外循环中学习的而不是通过人为先验手工制作的。 TTT的后续工作探索了诸如机器人操作[28]和运动[68]等应用这些应用通常需要对自监督任务进行不同的设计4.2.2 快速权重快速权重的一般思想是仅在最相关的数据上更新“快速”模型的参数而不是像传统做法那样在所有数据上更新“慢速”模型 [71]这个想法自20世纪80年代以来就存在了 [32]。 最相关的数据可以是测试实例本身因此TTT可以被视为快速权重的一种特殊情况先前关于快速权重的工作通常避免形成一个明确的学习问题来优化数据上的某些目标。 例如Hebbian学习和Hopfield网络的更新规则 [35] 只是简单地将 xxT或其某些变体 [4] 添加到每个输入 x的快速权重中。 相比之下TTT接受了明确制定学习问题的理念其中测试实例是泛化的目标。 TTT的更新规则也是一个明确的优化步骤快速权重程序员(FWPs)的想法是用一个“慢”模型来更新快速权重 [62]TTT的内循环权重 W可以被视为“快”的而外循环权重 θ则被视为“慢”的因此包含TTT层的网络可以被视为FWPs的一个特例 [43]类似于TTT可以被视为快速权重的一个特例。 上述使用Hebbian更新规则的FWP等价于线性注意力 [60]因此也等价于使用批量梯度下降的朴素TTT-LinearFWPs的定义非常广泛。 事实上所有具有某种门控机制的网络例如带有SwiGLU块的Transformer [63]也可以被视为FWPs的一个特例16。 最近的工作一直在尝试将FWPs用于语言建模Irie等人 [37] 设计了权重作为“慢”网络输出的“快”网络。Clark等人 [16] 给Transformer添加了一个快速权重的最终层其初始化被训练为慢权重TTT相对于现有FWP工作的贡献在于提出了一个明确的更新学习问题这使得能够借用诸如小批量和LN等学习工具4.3 Learning to Learn几十年来研究人员一直在争论学习如何学习也称为元学习或双层优化应该是智能的关键组成部分 [61, 6, 70, 47]。 在之前的工作中例如 [2]、[20] 和 [52]内循环每次从整个数据集而不是序列中学习因此外循环需要一组数据集或任务。 简而言之外循环是“比常规训练高一级”。由于很难收集数百万个数据集这个外循环很难扩展相比之下对于TTT每个序列本身就是一个数据集并定义了自己的泛化问题。内循环比常规训练“低一级”因此TTT的外循环只是监督学习经典问题的另一种解决方案而不是像跨数据集泛化那样的新问题设置如下表所示TTT的外循环与常规训练“处于同一级别”这使得TTT的外循环更容易扩展总之如上表所示TTT的论文将监督学习重新表述为学习如何学习具有两个嵌套循环外循环的高亮行与常规训练中的相同外循环的参数成为内循环的超参数。 直观地说内循环即TTT是常规训练的“下一级”最后如TTT论文所说为什么我们要研究TTT先前的工作通常尝试用机器学习来模拟人类学习其中训练是在一个洗牌后的i.i.d.实例数据集上进行的推理是在一个单独的测试集上进行的。然而人类并不是自然地通过i.i.d.实例学习也没有训练-测试分割人类学习与TTT(TTT的内部循环)有更有前途的联系其数据是一个潜在的非常长的序列具有强烈的时间依赖性任何数据片段都可以用于训练和测试。 这就是研究TTT的原因