学习（在测试时学习）：具有表现力丰富的隐藏状态的RNN

论文 | JAX代码库 | 环境设置 | 快速开始 | 推理基准

这是学习（在测试时学习）：具有表现力丰富的隐藏状态的RNN的官方PyTorch模型实现。 我们不建议使用此代码库进行训练，因为它是用纯PyTorch编写的，没有任何系统优化，所以训练速度会很慢，特别是当每个设备的批处理大小较小时。

对于训练代码，或复现我们论文中的结果，请查看我们的JAX代码库。对于推理内核，或复现我们论文中的速度基准，请查看我们的内核实现。

摘要

自注意力在长上下文中表现良好，但复杂度是二次方的。现有的RNN层具有线性复杂度，但它们在长上下文中的性能受限于隐藏状态的表达能力。我们提出了一类新的序列建模层，具有线性复杂度和表现力丰富的隐藏状态。核心思想是将隐藏状态本身设计为一个机器学习模型，而更新规则则是自监督学习的一个步骤。

由于隐藏状态甚至在测试序列上也通过训练进行更新，我们将这些层称为测试时训练（TTT）层。 我们考虑了两种实例化：TTT-Linear和TTT-MLP，它们的隐藏状态分别是一个线性模型和一个双层MLP。

环境设置

pip install "transformers[torch]"

快速开始

我们的实现基于Huggingface Transformers。您可以使用以下代码加载模型并生成文本。

from transformers import AutoTokenizer
from ttt import TTTForCausalLM, TTTConfig, TTT_STANDARD_CONFIGS

# 初始化TTT ttt-1b风格的配置
# configuration = TTTConfig(**TTT_STANDARD_CONFIGS['1b']) 等同于以下代码
configuration = TTTConfig()

# 从ttt-1b风格的配置初始化模型
model = TTTForCausalLM(configuration)
model.eval()

# 访问模型配置
configuration = model.config

# 分词器
tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-2-7b-hf')

# 预填充
input_ids = tokenizer("Greeting from TTT!", return_tensors="pt").input_ids
logits = model(input_ids=input_ids)
print(logits)

# 解码
out_ids = model.generate(input_ids=input_ids, max_length=50)
out_str = tokenizer.batch_decode(out_ids, skip_special_tokens=True)
print(out_str)

注意：这是为教学目的而设计的TTT层的简单实现。 这个模型可以使用Huggingface Accelerate或自定义训练循环进行训练。我们已经发布了更快的推理内核及其速度基准在此处。