arXiv:2407.04620v3 宣告类型: replace-cross 摘要：自注意力在长上下文情况下表现良好，但其复杂度呈平方级增长。现有的RNN层具有线性复杂度，但在长上下文情况下，其性能受限于隐藏状态的表达能力。我们提出了一种实用的框架，用于实例化具有线性复杂度和表达性强的隐藏状态的序列建模层。关键的想法是将隐藏状态本身视作一个机器学习模型，并将更新规则视为自我监督学习中的一步。由于隐藏状态即使在测试序列上也被用于训练，因此我们将这些层称为测试时训练（TTT）层。我们考虑了两种实例化方法：TTT-Linear和TTT-MLP，其中隐藏状态分别是一个线性模型和一个两层MLP。我们在参数量从125M到1.3B的范围内评估了我们的实现，并将其与一个强大的Transformer和现代的RNN——Mamba进行了比较。与Transformer类似，TTT-Linear和TTT-MLP可以通过对更多令牌进行条件处理来持续降低困惑度，而Mamba在16k上下文后无法继续改进。虽然TTT-MLP在内存I/O方面仍存在挑战，但它在长上下文方面展示了更大的潜力，这暗示了未来研究的一个有希望的方向。