lra-benchmark实战：如何评测注意力模型的长程能力

简介：该基准测试代码包基于谷歌研究院发布的远程竞技场数据集构建，面向需要评估模型长距离依赖建模能力的研究人员与开发者。项目使用PyTorch与HuggingFace Transformers库重新实现，替代了原本基于Jax或Flax的版本，代码抽象层次更直观，适合不熟悉JAX生态、希望快速上手远程竞技场任务的人群。压缩包共7个文件，其中4个Python脚本分别负责数据加载、配置解析与模型运行，另附说明文档、一键获取数据的Shell脚本及gitignore文件，整体体积仅5KB，结构十分精简。目前已有174人学习下载，可运行轻量级模型完成核心基准测试，并通过修改配置灵活替换数据集或模型结构；脚本采用相对路径执行，只需运行一条命令即可自动创建数据集目录，省去手动整理数据的麻烦；文档还对数据结构和实现要点作了说明，便于按需调整代码。对于想复现论文基线实验或扩展自定义任务的读者，这份资源提供了清晰且易于改写的起点。

1. lra-benchmark：它考的不是速度，而是长程注意力有没有"真本事"

lra-benchmark（Long Range Arena）是一套把"长程依赖"拆成六道考题的公开基准。它不测推理速度、不测参数量，只测一件事：当序列长到 1K、4K 甚至 16K token 时，你用的注意力机制还能不能抓住相距很远的信息点。

这个基准最有名的发现之一，是不少声称线性复杂度的注意力模型在前五个任务上成绩不错，一到 16K 的 Path-X 就集体翻车。所以它特别适合两类人：一类是正在设计新注意力算子的研究员，另一类是给业务模型做选型的工程师——在拿自己的数据上模型之前，先让它在 lra-benchmark 上交一次作业，比看任何自报的 FLOPs 都更能说明问题。本文按"任务拆解 → 本地跑通 → 接入自研模型 → 避坑 → 进阶验证"的顺序，把它讲透。

2. 六个任务拆解：Path-X 为什么让高效注意力集体翻车

2.1 一张表看懂六任务的输入形态与序列长度

LRA 设计六任务并不是随意拼的。它覆盖了文本、图像、逻辑表达式三种模态，并且把序列长度拉开了一个量级，从 1K 一直拉到 16K。下面这张表是官方设定口径，社区 fork 里会有小的预处理差异，但核心数字不变：

把这张表读透，就能理解为什么 LRA 后来成了高效注意力方向的标配实验：它把"局部性假设"和"全局性假设"同时塞进一套评测里。Text 和 Retrieval 考的是语料级别的语义，Image 和 Pathfinder 考的是容易被局部卷积偷走的视觉能力，而 Path-X 直接把图像放大到 16K，逼着你证明自己的注意力真的在往远处看。

2.2 线性复杂度模型为什么会在这套题上漏底

理论直觉往往和实际成绩有出入。线性注意力把 softmax 分解成"特征映射内积"，运算复杂度降到了线性，但代价是注意力分布被限制在低秩空间里。低秩对"局部模式"足够用，对"两个特定点之间长达几千 token 的精确关系"就很吃力，这是它翻车的数学根源。

Pathfinder 这类任务看起来简单，实际上是视觉版的"找路径"：两个点之间有没有连通路径，模型必须同时记住多个位置的像素信息再做空间推理。Path-X 把这个问题推到 16K 之后，稀疏注意力模型常用的"局部窗口 + 少量全局 token"策略也不够用了——全局 token 数量有限，覆盖不了任意两点之间所有可能的配对。这就是公开结果里最扎眼的现象：全注意力的 dense Transformer 在 Path-X 上仍然很强，而大量高效注意力变体只剩 50%~60% 的准确率，跟瞎猜差不多。

这也是 lra-benchmark 最有价值的地方：它专门把"效率换能力"的账本翻出来给你看。你省了多少复杂度，在长程题上要还多少能力债，六个任务跑完一对比就一目了然。别以为 Path-X 是视觉任务就跟 NLP 无关——长文档问答、长视频、多模态检索里，"超长距离精确关联"恰恰是最难攻的一块，提前暴露问题不是坏事。

2.3 选基线的常用做法：上界、下界和参考点

第一次跑 lra-benchmark，建议先选三条基线：dense Transformer 作为上界，一个你了解底细的线性注意力模型（比如基于随机特征的 kernel 方法）作为下界，再找一个稀疏窗口类模型作为中间参考。这样你的自研模型跑出来的分数就有坐标系了：它必须比下界高一截、并且逼近上界才有意义。

从公开榜单的普遍结论看，BigBird 这类把局部窗口、全局 token 和随机稀疏结合起来的模式，是高效模型里综合成绩最稳的之一；而纯随机特征类的线性注意力在 Path-X 上几乎必然掉队。这个排序在不同 fork 里会有一两个点的浮动，但你只需要记住：任何新注意力接入后，第一件事不是看绝对分数，而是看它相对这三条基线的差距。这个差距，就是你在长程能力上的真实位置。

3. 本地跑通 lra-benchmark：环境、目录与最小命令

3.1 克隆、建环境、装依赖

第一步是拿到一份能跑的代码。无论你克隆的是上游的 long-range-arena 还是社区维护的 lra-benchmark fork，目录骨架基本一致：

# 用你实际拉取的 fork 地址替换占位符；下文统一把仓库根目录叫 lra-benchmark git clone <remote-url> lra-benchmark cd lra-benchmark # 建议用独立虚拟环境，避免把全局环境搞脏 python3 -m venv .venv source .venv/bin/activate # 依赖以仓库里的 requirements.txt 为准，不要自己另配版本 pip install -r requirements.txt

这里有个常见选择要说明：原版代码是 TensorFlow 时代的技术栈写的，现在社区里大部分可用的 lra-benchmark fork 都是 PyTorch 移植版，依赖通常只有 torch、numpy、yaml 几个核心包。如果你拿到的是原版 TF 代码，装环境时大概率会遇到 TF 与 Python 版本互斥的问题。我的建议是直接找活跃维护的 PyTorch 移植版，省下的时间足够多跑两个任务。

3.2 先看懂目录：任务在 lra/ 下，配置在 configs 里

lra-benchmark/ ├── bash_scripts/ # 一键跑全量任务的脚本集合 ├── lra/ │ ├── common/ # runner 共用工具：指标、种子、模型基类 │ ├── listops/ # 任务目录：含 runner.py 与 configs/ │ ├── text/ │ ├── retrieval/ │ ├── image/ │ ├── pathfinder/ │ └── path_x/ # 有些 fork 叫 pathfinder_x，先 ls 确认 └── requirements.txt

每个任务目录里最核心的是 runner.py 和 configs/ 下的 yaml 配置。runner.py 负责读配置、加载数据、训练、评估；yaml 里定义了学习率、batch size、层数、注意力类型等全部超参数。任务之间的差异被封装在各自目录里，所以换任务的成本很低，通常只是换 config 路径。

3.3 先跑冒烟：Image 或 ListOps，100 步验证管道

直接上全量任务是最常见的翻车方式。正确做法是先跑一个 100 步的冒烟测试，确认数据加载、模型前向、loss 回传整条链路是通的：

# --max-steps 100 只走 100 步训练，验证整条 pipeline python lra/image/runner.py \ --config lra/image/configs/image_default.yaml \ --max-steps 100 \ --log-every-n-steps 10

参数说明：--config 指定该任务的 yaml；--max-steps 是训练总步数，冒烟时给 100；--log-every-n-steps 控制日志频率，10 步打一次 loss。如果 100 步内 loss 从初始值明显下降，说明数据管道和模型梯度都是通的，可以放心放量。若 loss 是 NaN 或者不降，先别调参，回到数据预处理和模型 forward 里查，这一步最省时间。

3.4 全量六任务的命令与预算

跑通冒烟后，按任务从小到大的顺序逐个放量。常见做法是写一个循环脚本，每个任务用独立 seed 跑三遍，取平均作为该模型在 LRA 上的成绩：

for task in listops text retrieval image pathfinder path_x; do for seed in 0 1 2; do python lra/$task/runner.py \ --config lra/$task/configs/${task}_default.yaml \ --seed $seed \ --max-steps 100000 done done

逻辑说明：外层循环遍历六个任务，内层循环跑三个 seed。注意上面循环里的配置文件名是常见命名，不同 fork 可能把 Retrieval 的配置叫别的名字，跑之前先 ls 一下 configs 目录。100000 步是官方参考量级，实际收敛大多在 5 万步以内；如果你只做对比实验、不想等太久，可以先用 3 万步跑一轮看趋势。

预算上，单张 GPU 跑完六个任务 × 3 seed，按模型不同大约需要一到三天，Path-X 最费时间。所以时间敏感的项目建议先跑 Image、ListOps、Text 三个任务出初步结论，确认模型有价值后再补 Pathfinder 和 Path-X。

配置里需要认识的几个关键参数，基本长这样：

4. 把自研注意力接进 lra-benchmark：接口、注册与自检

4.1 三种输入形态决定模型要先分清"喂进来的是什么"

LRA 任务的输入不是统一的词汇表，而是分三种形态：

• ListOps / Text：离散 token，需要模型自己建 embedding；
• Image / Pathfinder / Path-X：连续像素特征，每个 token 是一个 1D 特征向量，没有 embedding 环节；
• Retrieval：两篇文本拼接，处理方式仍然是 token 序列。

如果你只按 NLP 的习惯写 forward，拿到 image 任务就会在维度上直接报错。我在接入时一般会在模型开头打一行形状断言，把"这个任务到底喂什么"变成一次显式检查，而不是黑盒式调参：

# 在 forward 最前面断言输入形状，任务形态错了立刻报错 assert inputs["inputs"].dim() >= 2, "输入至少是 [batch, length]" if inputs.get("input_ids") is not None: x = self.embed(inputs["input_ids"]) # 离散任务：token 转 embedding else: x = inputs["inputs"] # 连续任务：像素特征直接进编码器

这个分支判断是接入 LRA 最容易写错的地方，因为不同任务的 runner 传给模型的 dict 结构不一样。动手写模型前，先打印一次实际收到的 inputs，看清里面有 input_ids、inputs、mask、target 这些 key 的哪些组合。

4.2 实现 Model 子类：forward 之外还要管好 train/eval

各任务 runner 通过一个基类接口驱动模型，最省心的做法是继承任务目录里的模型基类，只实现 forward：

class MyModel(BaseModel): def __init__(self, config): super().__init__(config) self.encoder = MyEncoder(config) self.head = nn.Linear(config["hidden_size"], config["num_classes"]) def forward(self, inputs, training=True): # training=True 时走 dropout；eval 时关闭，防止评估分数偏低 mask = inputs["mask"] x = inputs["input_ids"] if "input_ids" in inputs else inputs["inputs"] logits = self.head(self.encoder(x, mask)) return {"logits": logits}

这里说两个重点。第一，eval 和 train 的差别要自己在 forward 里用 training 开关处理：很多移植版 runner 会调 model.eval()，但如果你在 forward 里硬编码了 dropout，评估分数就会比真实水平低一截，这是"结果玄学"的常见来源。第二，logits 的输出结构要和 runner 的约定一致，多数实现期望一个包含 logits 的 dict，而不是裸张量，写之前去看一眼现有 dense 模型的返回格式。

4.3 把模型注册进 registry，用 --model 指定

runner 通过 --model 参数在注册表里查模型，所以接入最后一步是把模型类挂进去：

MODEL_REGISTRY = { "transformer": DenseTransformer, # 已有的 dense 基线，别删 "my_model": MyModel, # 你的新注意力模型 }

然后跑任务时指定模型名：

python lra/pathfinder/runner.py \ --config lra/pathfinder/configs/pathfinder_default.yaml \ --model my_model

注意：config 文件里的 attention 类型参数和 --model 是两层的东西。--model 决定实例化哪个类，config 里的参数只是喂给这个类的超参。很多人把两者搞混，结果 self.encoder 内部还在按 config 里的旧 attn_type 分支选择实现，新模型根本没有被用到。

4.4 接入后的自检：20 步内看三件事

模型接入后，别急着上全量。先用 20 步观察三个信号：

• logits 的形状是否与 num_classes 一致；
• loss 是否随步数单调下降；
• eval 准确率是否明显大于随机猜测。

我习惯写一个 20 步冒烟脚本，把这三个信号全打出来，确认无异常再放量。任何一步不对都要回到 forward 里查，而不是调学习率——在 20 步尺度上，学习率不是首因。跑完冒烟再按第 3 章的循环脚本上全量，基本不会浪费算力。

5. 跑 lra-benchmark 的避坑清单：从训练不收敛到结果不可复现

5.1 训练不收敛的两个坑：先怀疑配置，再怀疑模型

坑一：Path-X 的 loss 纹丝不动，准确率钉在 50%。现象是训练几千步了，验证准确率一直停在随机猜测附近。原因通常不是模型坏了，而是这个任务序列长、训练慢，且很多移植版在转换配置时把学习率放大了一两个数量级；另外如果开了混合精度，长序列上的数值溢出会让梯度变成 NaN。解决：把 config 里的 learning_rate 改回 1e-4 量级，先关闭混合精度，确认 loss 一开始就在下降；还是不行就把序列截断到 1K 验证代码逻辑，再放回全长度。这个排查顺序能避免你在错误的层上调半天。

坑二：小任务正常，一换任务 loss 就起飞。现象是 ListOps 跑得好好的，切到 Image 后 loss 直接发散。原因多半是输入形态变了，embedding 分支还在被错误调用，连续像素被当成离散 id 过了一个大 embedding 表。解决：在 forward 里按 4.1 的方式断言输入形态，打印一次 inputs 的 key 和 dtype，确认走的是连续特征分支还是离散 token 分支。

5.2 复现与评估相关的三个坑：种子、切分和显存

坑三：同一个模型跑两次，准确率差两三个点。这是 lra-benchmark 下最常见的"结果玄学"。原因是多进程数据预处理和装桶打乱时，随机种子没有在数据洗牌之前固定，导致每次喂给模型的样本顺序都不同。解决：把模型框架的种子、numpy 种子、Python hash 种子全部在 runner 入口固定，并且每个实验用 3 个不同 seed 各跑一遍，报告均值而不是单次成绩。别信单次跑出来的最高分，那是自欺欺人。

坑四：eval 分数比训练还离谱，甚至超过 dense 上界。现象是验证集准确率接近 100%，而训练集表现正常。原因多数是评估阶段没有关 dropout，或者把训练集当验证集喂了进去；还有一种可能是 Retrieval 任务的切分没对齐，模型在验证时见到了训练样本。解决：检查 runner 里 eval 时是否调用了 model.eval()，核对数据切分函数，并抽样验证训练集和验证集没有交集。这个坑最容易发生在你自己改了数据加载逻辑之后，改动越少越安全。

坑五：显存直接爆掉。现象是 Image 能跑，切到 Path-X 就 OOM。原因是 Path-X 是 16K 序列，任何平方复杂度的中间张量都会把显存吃穿，而且不少人会顺手把 batch size 调大。解决：Path-X 上 batch size 固定为 16 或更小，开启 gradient checkpointing，把不需要梯度的大中间量主动 detach 掉。如果你在跑自研模型，先用 Pathfinder 的 1K 序列验证正确性，再上 Path-X——这样爆显存时你能确定问题在显存而不在模型。

6. 进阶验证：用长度消融证明你的模型真的在长程上变强了

六任务分数只能说明"总分不差"，不能证明你的模型赢在长程能力上。一个只在局部卷积上更强的模型，完全可能在 Pathfinder 上得分更高。要戳穿这个假象，做一次长度消融：以 ListOps 为例，按输入序列长度分桶统计准确率，看你的模型是"越长越弱"还是"全程稳定"。

# 按输入长度分桶统计准确率，验证"越长越稳"还是"越长越崩" import numpy as np lengths = [] # 每个样本的输入长度，从 runner 的日志里收集 correct = [] # 每个样本是否正确 edges = [0, 512, 1024, 2048, 4096] for lo, hi in zip(edges[:-1], edges[1:]): idx = (np.array(lengths) >= lo) & (np.array(lengths) < hi) print(f"len {lo}-{hi}: acc={np.mean(correct[idx]):.3f}")

判断标准很直接：dense Transformer 随着长度增长会缓慢下降，线性注意力模型会在最长桶上剧烈跳水；如果你的模型在长桶上保持平缓，说明它是真的在利用长程信息，而不是靠局部模式硬撑。报告结果时，把三组 seed 的均值±标准差写上，和 dense 基线在同一训练预算下对比，别拿自己调参后的模型去碰别人没调参的基线。

我现在的习惯是：任何新注意力模块想进业务，先花一个晚上把 lra-benchmark 的 1K 任务跑通，再花两三天补上 Path-X；如果 Path-X 撑不住，绝不用"我们的场景没这么长"来安慰自己。跑完 LRA 再回头看业务数据，你会对"长程"这两个字有完全不同的敬畏。希望帮到你。

本文还有配套的精品资源，点击获取