Nyströmformer 源码解析与实战指南：基于 Nyström 采样的 O(n) 自注意力近似与 Transformers 实现

导读

Nyströmformer 是 2022-01-11 合入 Hugging Face Transformers（论文发布于 2021-02-07）的一个高效 Transformer 架构：它把经典的 Nyström 方法（低秩矩阵近似算法）迁移到自注意力上，通过少量"地标点（landmarks）"与迭代式 Moore-Penrose 伪逆把 softmax 注意力从二次复杂度降到 O(n)，从而可以处理数千 token 的长序列。本文以仓库内模型文档 docs/source/en/model_doc/nystromformer.md 为主体，结合核心实现 modeling_nystromformer.py、配置类 configuration_nystromformer.py 与官方 checkpoint 转换脚本，完整讲解其数学原理、逐参数配置、六大公开 API（模型 + 五个下游任务头）及加载/转换/使用细节。

---

1. 模型概述：用低秩近似对抗自注意力的平方复杂度

1.1 要解决的问题

标准 Transformer 的自注意力需要计算 n×n 的注意力矩阵（n 为序列长度），时间与空间复杂度均为 O(n²)。这一瓶颈限制了 Transformer 在长序列任务（数千 token）上的应用。Nyströmformer 通过改造 softmax 注意力矩阵的计算方式来突破这一限制，论文发表于 arXiv 2102.03902（Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention，作者 Yunyang Xiong、Zhanpeng Zeng 等）。

1.2 核心思路

我的思路基于将 Nyström 方法适配到标准自注意力的近似上，实现 O(n) 复杂度。Nyströmformer 的扩展性使其能够应用于包含数千 token 的更长序列。——论文摘要（仓库内模型文档引用）

论文在 GLUE 与 IMDB 等标准长度任务上报告了与标准自注意力相当、个别任务略优的效果；在 Long Range Arena（LRA）长序列基准上，相对其它高效注意力方法表现良好。模型在仓库内的源码实现与配置均严格对应这一设计，仓库不包含独立的推理/训练脚本，全部能力经由 Transformers 标准 API 暴露（见第 5 节）。

Nyströmformer 在本仓库中的模型贡献者与社区结构：

• 由社区开发者 novice03 贡献，原始代码来自 mlpen/Nystromformer（论文作者方仓库，仅作背景参考，不在本仓库内）；
• 模型类别在 configuration_nystromformer.py 中注册为 model_type = "nystromformer"；
• 官方 checkpoint 为 uw-madison/nystromformer-512，该名称在配置类装饰器（configuration_nystromformer.py）与转换脚本中作为默认示例 checkpoint 出现。

1.3 顶层架构

从 modeling_nystromformer.py 的整体结构看（共 946 行），模块组织如下：

组件类	作用
NystromformerEmbeddings	词嵌入 + 位置嵌入 + 段落类型（token_type）嵌入，Bert 风格（见下方"与 BERT 的关系"）
NystromformerSelfAttention	全模型核心：Nyström 近似自注意力（landmark 采样、三个 softmax 核、迭代伪逆、深度卷积）
NystromformerSelfOutput / NystromformerAttention	自注意力 + 残差 + LayerNorm 输出（对齐 BERT 结构）
NystromformerIntermediate / NystromformerOutput	两层 MLP + 残差 + LayerNorm（FFN 块）
NystromformerLayer / NystromformerEncoder	单层 Transformer Layer 与 12 层编码器堆叠，输出隐藏状态/注意力权重，支持梯度检查点
NystromformerPreTrainedModel	基类，supports_gradient_checkpointing = True，前缀 nystromformer
NystromformerModel + 5 个任务头	裸编码器与 MLM/分类/多选/Token 分类/问答

该文件大量复用 BERT 代码，多个模块标注 # Copied from transformers.models.bert.modeling_bert.BertXXX（如 SelfOutput、Intermediate、Output、预测头等），说明其整体骨架沿袭 BERT——唯一的本质差异集中在自注意力层，这也是理解本模型的关键。

1.4 嵌入层细节

NystromformerEmbeddings（modeling_nystromformer.py）有两点值得注意：

• 位置嵌入维度是 max_position_embeddings + 2，且位置 ID 从 2 开始递增（torch.arange(config.max_position_embeddings).expand((1, -1)) + 2），这与 checkpoint 训练时在序列前保留特殊 token 的偏移约定有关，使用官方 checkpoint 时无需自行处理；
• position_ids 与 token_type_ids 以 nn.Buffer（persistent=False）注册，便于模型 trace/导出时自动补齐，且 _init_weights 会对其做显式初始化（modeling_nystromformer.py）。

---

2. 核心原理：Nyströmformer 自注意力是如何计算的

这一节是理解该模型的重中之重，全部推理均可在 NystromformerSelfAttention（modeling_nystromformer.py）的 forward 中找到一一对应代码。

2.1 标准注意力回顾与替代方案设计

标准 softmax 注意力：Attention(Q, K, V) = softmax(QKᵀ/√d) V，瓶颈在 QKᵀ 形成的 n×n 矩阵。Nyströmformer 的替代计算分两条路径：

路径 A：全注意力兜底。当 num_landmarks == segment_means_seq_len（默认均为 64，二者相等）时，模型在 forward 中走"标准"分支：直接计算完整 QKᵀ、softmax、乘以 V。这意味着默认配置下该分支会退化为标准注意力语义（此时 seq_len 实际接近 64），等效于对短序列用精确注意力、长序列用近似注意力的自适应行为——源码注释亦明示：当 landmark 数与 segment 均值序列长度相等时不做采样近似。

路径 B：Nyström 近似（长序列主路径），即当两者不等时执行。整条流水线如下。

2.2 缩放：把 √d 拆成两半

在 forward 开头（L168-L169）：

query_layer = query_layer / math.sqrt(math.sqrt(self.attention_head_size))
key_layer = key_layer / math.sqrt(math.sqrt(self.attention_head_size))

即将 1/√d 均分到 Q 与 K 两侧，配合后续所有矩阵乘积，数学上仍等价于整体缩放 1/√d，但避免了在计算中单独维护一个大矩阵的额外算子。Q、K、V 由三个独立 nn.Linear 投影得到，并按头数切分为 (batch, heads, seq, head_dim)。

2.3 Landmark（地标点）采样：segment-means 下采样

对 Q 与 K 各取少量地标点。源码实现采用 segment-means（段均值）策略（L182-L195）：

q_landmarks = query_layer.reshape(
    -1, self.num_attention_heads, self.num_landmarks,
    self.seq_len // self.num_landmarks, self.attention_head_size,
).mean(dim=-2)
k_landmarks = key_layer.reshape(...).mean(dim=-2)

序列被等分成 num_landmarks 段，每段长度 seq_len // num_landmarks，段内向量求均值即得到地标点。因此要求 seq_len 能被 num_landmarks 整除。值得注意的是：源码中的 self.seq_len 取自 config.segment_means_seq_len（默认 64），而非运行时实际序列长度，此处的"段均值"始终在"名义序列长度 64"上进行——这意味着真正支撑长序列的近似能力依赖在 64 之外扩展位置嵌入与上下文长度，而注意力核计算在采样的 64×64 规模上完成，这正是复杂度从 O(n²) 降为 O(n) 的关键。

2.4 三个 softmax 核重建注意力矩阵

设 Q̄、K̄ 为地标点矩阵，S 为地标-键注意力（用于吸收 key 的归一化），Nyströmformer 用一个低秩乘积 A ≈ K₁ · (K₂)⁺ · K₃ 近似原注意力矩阵，源码（L197-L209）：

kernel_1 = torch.nn.functional.softmax(torch.matmul(query_layer, k_landmarks.transpose(-1, -2)), dim=-1)   # Q 对 K̄
kernel_2 = torch.nn.functional.softmax(torch.matmul(q_landmarks, k_landmarks.transpose(-1, -2)), dim=-1)   # Q̄ 对 K̄
attention_scores = torch.matmul(q_landmarks, key_layer.transpose(-1, -2))                                  # Q̄ 对 K
# ...（若 attention_mask 不为空则 attention_scores += attention_mask）
kernel_3 = nn.functional.softmax(attention_scores, dim=-1)

attention_probs   = torch.matmul(kernel_1, self.iterative_inv(kernel_2))   # n×m · m×m
new_value_layer   = torch.matmul(kernel_3, value_layer)                    # m×n · n×d
context_layer     = torch.matmul(attention_probs, new_value_layer)         # n×m · m×d

数学要点：kernel_2 是 m×m 的（m = num_landmarks = 64），对它求 Moore-Penrose 伪逆 (kernel_2)⁺；attention_probs 的"行"即近似后的注意力权重，最大规模被钳制在 n×m，替代了原始 n×n，从而获得线性复杂度。这里的 kernel_2⁺ 正是 Nyström 方法中的矩阵求逆步骤。

2.5 迭代式 Moore-Penrose 伪逆

kernel_2 是 softmax 输出（逐行和为 1 的正矩阵），其伪逆通过 iterative_inv（L139-L159）迭代计算，默认 6 次迭代，每次更新公式为：

value = 0.25 · value · (13I − key_value·(15I − key_value·(7I − key_value)))

其中 key_value = kernel_2 · value。该迭代式源自论文采用的逆矩阵三次多项式加速方案，把伪逆计算从 O(m³) 的 SVD 简化为矩阵乘法迭代。初始化 value 有两种策略，由配置项 inv_coeff_init_option 控制（详情见第 3 节参数表），源码注释说明：

• "original"（默认）：系数取 1 / max(Σ_k key)，实现更保守；
• 精确系数版本：取 1 / ||key||₁ 按行归一，收敛更快。

2.6 深度卷积跳跃连接

近似注意力还可能损失细节，因此实现加入一个可学习的补偿项（L129-L137 与 L211-L212）：

self.conv = nn.Conv2d(
    in_channels=self.num_attention_heads, out_channels=self.num_attention_heads,
    kernel_size=(self.conv_kernel_size, 1), padding=(self.conv_kernel_size // 2, 0),
    bias=False, groups=self.num_attention_heads,
)
...
context_layer += self.conv(value_layer)

它是在各注意力头上独立运行的 depthwise 卷积（groups=heads），沿序列维度（kernel 宽为 1，即只沿长度方向滑窗）作用在 value_layer 上，输出直接与上下文结果相加，帮助重建局部邻域信息。注意若 conv_kernel_size=None 则完全不创建该卷积层（见配置节）。

2.7 输出组装

context_layer 由 (batch, heads, seq, head_dim) 转回 (batch, seq, all_head_size)（L214-L216），随后进入与 BERT 相同的 SelfOutput（残差 + LayerNorm）、FFN 块完成单层；NystromformerEncoder 按 num_hidden_layers 循环堆叠并支持返回隐藏状态与注意力矩阵。output_attentions=True 时返回的是 m×n 规模（或退化为全注意力的 n×n）的近似注意力权重，不再是标准 n×n 矩阵——如果你用 output_attentions 做可解释性分析，需要注意这一点。

---

3. NystromformerConfig：Nyström 专属参数与全量配置表

配置类定义于 configuration_nystromformer.py。它继承 PreTrainedConfig、model_type = "nystromformer"，并在构造时严格校验字段类型。

3.1 Nyström 专属参数（本文档核心配置项）

参数	类型	默认值	含义与约束
num_landmarks	int	64	Nyström 近似 softmax 注意力矩阵时使用的地标点（landmark）数量；等价于 Nyström 采样点个数 m
segment_means_seq_len	int	64	segment-means 下采样使用的（名义）序列长度。要求能被 num_landmarks 整除（代码中按 seq_len // num_landmarks 分组取均值）；当 segment_means_seq_len == num_landmarks 时退化为精确全注意力路径
conv_kernel_size	int	65	自注意力内 depthwise 卷积的核大小（沿序列维）。设为 None 可关闭该卷积补偿层；padding 由源码自动按 kernel_size // 2 计算，保证长度不变
inv_coeff_init_option	bool	False	是否对迭代法求 Moore-Penrose 伪逆的初值 Z₀ 使用"精确系数"计算（`1/

顺带说明：NystromformerSelfAttention 中有一处历史遗留逻辑，当 inv_coeff_init_option=True 时尝试读取 config["inv_init_coeff_option"]（L118-L121）。从配置类看当前公开字段名为 inv_coeff_init_option，实际使用中应直接设置该公开字段；这一内部分支不影响默认 False 路径，提请注意以免配置名混淆。

3.2 通用超参（默认 BERT-base 规模）

以下参数与 Bert 同构，默认值直接内嵌于类定义（configuration_nystromformer.py）：

参数	默认值	说明
vocab_size	30000	词表大小
hidden_size	768	隐藏维度（须能被 num_attention_heads 整除，否则构造时报错）
num_hidden_layers	12	编码器层数
num_attention_heads	12	注意力头数
intermediate_size	3072	FFN 中间维度
hidden_act	"gelu_new"	激活函数（映射自 ACT2FN）
hidden_dropout_prob / attention_probs_dropout_prob	0.1	隐层 / 注意力 Dropout
max_position_embeddings	510	最大位置嵌入数（配合位置偏移从 2 开始的设计）
type_vocab_size	2	段落类型数（如 sentence A/B）
initializer_range	0.02	参数初始化范围
layer_norm_eps	1e-5	LayerNorm epsilon
pad_token_id / bos_token_id / eos_token_id	1 / 0 / 2	特殊 token ID
add_cross_attention	False	是否启用交叉注意力（encoder-decoder 场景）
tie_word_embeddings	True	是否共享输入/输出词嵌入（MLM 头）

3.3 用法示例

from transformers import NystromformerModel, NystromformerConfig

# 初始化一个 uw-madison/nystromformer-512 风格的配置（全部取默认值）
configuration = NystromformerConfig()

# 由该配置随机初始化模型
model = NystromformerModel(configuration)

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

# 长序列适配：例如把位置上限从 510 扩到 4096、并调整 landmark 策略
long_config = NystromformerConfig(
    max_position_embeddings=4096,
    num_landmarks=256,          # 更多 landmark 提升近似精度
    segment_means_seq_len=256,  # 与 num_landmarks 相等时走精确注意力；不相等时走 Nyström 近似
    conv_kernel_size=65,
    inv_coeff_init_option=True, # 使用精确伪逆初值加速收敛
)

---

4. 六大公开 API 逐一解析

模型文档按 [[autodoc]] 声明以下对象，对应的类实现均在 modeling_nystromformer.py 与 init.py 中导出。这些类同时挂载在 Auto 体系下（见 modeling_auto.py 中的 AutoModelForMaskedLM 等映射），可无缝接入 Trainer / pipeline / 任务脚本。注意：这些类的 forward 文档由仓库内 auto_docstring 机制自动生成（auto_docstring.py），标准入参（input_ids、attention_mask、token_type_ids、position_ids、inputs_embeds、output_attentions、output_hidden_states、return_dict）含义与 BERT 家族一致，本文按任务差异给出每类入参要点。

4.1 NystromformerModel

裸 Transformer 编码器。入参与第 2.7 节一致；返回 BaseModelOutputWithPastAndCrossAttentions：last_hidden_state、hidden_states（开启 output_hidden_states 时）、attentions（开启 output_attentions 时为近似注意力）、cross_attentions。典型用途是取序列表示做下游自定义头。

4.2 NystromformerForMaskedLM

裸编码器 + NystromformerOnlyMLMHead（BERT 风格的预测头，L505-L577）。

• 入参 labels：形状 (batch_size, sequence_length) 的整数，取值 [-100, 0, ..., vocab_size-1]；-100 的 token 不参与损失计算（即被 mask 的位置以外的 token 置为 -100）；
• 返回 MaskedLMOutput（loss/logits/hidden_states/attentions）；
• 输入输出词嵌入默认共享（tie_word_embeddings=True），并通过 _tied_weights_keys 声明绑定关系。

4.3 NystromformerForSequenceClassification

编码器 + NystromformerClassificationHead（L580-L684）：取序列首 token features[:, 0, :]（等价 [CLS]），经 dropout→dense→激活→dropout→out_proj 输出 logits。

• 入参 labels：形状 (batch_size,)；num_labels == 1 时走 MSE 回归损失，num_labels > 1 且标签为 long/int 时走单标签 CE，多标签时走 BCEWithLogits（problem_type 可自动推断或手工指定）；适用于 GLUE 类任务；
• 返回 SequenceClassifierOutput。

4.4 NystromformerForMultipleChoice

编码器 + 两层分类头（L687-L789）。

• 入参形状注意：input_ids 形状为 (batch_size, num_choices, sequence_length)，模型内部先拍平为 (-1, seq_len) 过编码器，取首 token 池化后经 pre_classifier + ReLU + classifier 得到每选项 logit，再 reshape 回 (batch_size, num_choices)；
• labels 形状 (batch_size,)，取值 [0, ..., num_choices-1]；
• 返回 MultipleChoiceModelOutput（含最终 logits 与可选损失）。

4.5 NystromformerForTokenClassification

编码器 + 对每个 token 的线性分类头（L792-L855）。

• labels 形状 (batch_size, sequence_length)，取值 [0, ..., num_labels-1]，整序列 CE 损失；
• 返回 TokenClassifierOutput；适用于命名实体识别等序列标注任务。

4.6 NystromformerForQuestionAnswering

编码器 + 抽取式问答头（L858-L934）：构造时将 num_labels 固定为 2，输出 start_logits 与 end_logits。

• 入参 start_positions / end_positions：形状 (batch_size,)；越界位置会被 clamp 到序列末尾、并以 ignore_index 屏蔽后取 CE，最终损失为起止两端损失均值；
• 返回 QuestionAnsweringModelOutput。

---

5. 实际使用指南：加载、微调与转换

5.1 通过 Auto API 快速加载（推荐）

Nyströmformer 在 Auto 体系中注册齐全：配置映射见 auto_mappings.py，模型与任务头映射见 modeling_auto.py 及后续各任务类映射。加载官方 checkpoint：

from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoConfig

model_name = "uw-madison/nystromformer-512"   # 官方示例 checkpoint
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)

inputs = tokenizer("Transformers are <mask>.", return_tensors="pt")
outputs = model(**inputs, labels=inputs["input_ids"])   # 传 labels 即得 MLM loss

分词：Nyströmformer 自身不携带分词器实现，Auto 体系按 tokenization_auto.py 的映射复用 BERT 家族的分词后端（checkpoint 发布时即配合 BERT 风格 tokenizer），因此加载官方 checkpoint 时用 AutoTokenizer 即可获得与之匹配的分词结果。

如需 AutoModel / 分类 / QA 等：

AutoModel.from_pretrained(model_name)                       # NystromformerModel
AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
AutoModelForTokenClassification.from_pretrained(model_name, num_labels=...)
AutoModelForMultipleChoice.from_pretrained(model_name, num_choices=...)
AutoModelForQuestionAnswering.from_pretrained(model_name)

5.2 从论文原始仓库 checkpoint 转换（进阶）

若你拥有 mlpen/Nystromformer 训练得到的 PyTorch checkpoint，可通过仓库内脚本 convert_nystromformer_original_pytorch_checkpoint_to_pytorch.py 转成 Transformers 格式。脚本要点：

• 按规则重命名键：如 model. 前缀去除、transformer_{N} → encoder.layer.{N}、mha.attn → attention.self、W_q/W_k/W_v → self.query/key/value、norm1/norm2 → 对应 LayerNorm、ff1/ff2 → FFN 的 dense、mlm 头 → cls.predictions 等；丢弃 pooler / 分类器 / conv.bias 权重；
• 补写 cls.predictions.bias 与 position_ids 缓冲（torch.arange(max_position_embeddings)+2，与第 1.4 节偏移约定呼应）；
• 命令行用法：

python src/transformers/models/nystromformer/convert_nystromformer_original_pytorch_checkpoint_to_pytorch.py \
    --pytorch_model_path /path/to/original/checkpoint.pt \
    --config_file /path/to/config.json \
    --pytorch_dump_path ./nystromformer-hf-output

5.3 长序列使用建议与限制

• 默认配置下模型的位置上限为 510（max_position_embeddings），如需处理数千 token，务必同步扩 max_position_embeddings（自行扩展时需理解 embedding 初始化行为，与官方 checkpoint 不完全等价）；
• num_landmarks 与 segment_means_seq_len 的组合决定了近似精度与复杂度，二者相等即退化为精确注意力（见 2.1）；扩大 num_landmarks 可提高近似保真度但会增加 O(n·m) 成本；
• 修改 conv_kernel_size 为 None 可去掉深度卷积补偿分支，用于对比实验或减小开销；
• 若开启 output_attentions=True，返回的注意力为低秩近似结果（m×n 规模），尺寸与语义和标准注意力不同，下游分析需适配。

---

6. 任务指南速查与进一步阅读

本文档在 "Resources" 中列出 Nyströmformer 可直接接入的标准任务指南，转换为仓库根目录相对路径后如下（每个指南均含配套示例脚本与 pipeline 用法）：

• 文本分类（sequence classification）
• Token 分类（token classification）
• 问答（question answering）
• 掩码语言建模（masked language modeling）
• 多项选择（multiple choice）

这些任务分别与第 4 节中的五个带任务头类一一对应，可直接套用示例脚本的 --model_name_or_path 参数指向 uw-madison/nystromformer-512（注意把默认 --max_seq_length 控制在模型位置上限内）。

实现验证资料：

• 模型实现：src/transformers/models/nystromformer/modeling_nystromformer.py
• 模型配置：src/transformers/models/nystromformer/configuration_nystromformer.py
• 权重转换脚本：src/transformers/models/nystromformer/convert_nystromformer_original_pytorch_checkpoint_to_pytorch.py
• 模块导出：src/transformers/models/nystromformer/init.py
• 测试覆盖：普通单测与集成测试见 tests/models/nystromformer/test_modeling_nystromformer.py，其中模型测试类同时继承 ModelTesterMixin 与 PipelineTesterMixin，覆盖前向、梯度检查点、任务头与 pipeline 兼容性；集成测试则校验真实 checkpoint 行为。

7. 小结

Nyströmformer 在本仓库中的实现完整体现了论文的三要素：segment-means 地标采样把注意力核矩阵压缩到 m×m 规模、迭代式 Moore-Penrose 伪逆（6 次迭代、两种初值策略）以 O(m³) 的矩阵乘近似替代精确求逆、depthwise 卷积跳跃连接补偿低秩近似的局部信息损失，从而把自注意力整体推到 O(n) 复杂度。结合 NystromformerConfig 的四项专属参数，开发者可以在"精确注意力"与"线性近似"之间连续调节，并以 Auto API 或官方 checkpoint 转换脚本快速落地到文本分类、序列标注、抽取式问答、多项选择与掩码语言建模等标准任务上——是一个理解"高效注意力变体如何在框架内落地"的理想参考实现。