Fairseq 中 Linformer 线性复杂度自注意力在 RoBERTa 预训练中的配置与源码级实现解析【免费下载链接】fairseqFacebook AI Research Sequence-to-Sequence Toolkit written in Python.项目地址: https://gitcode.com/gh_mirrors/fa/fairseq导读本指南围绕 fairseq 仓库examples/linformer/提供的 Linformer 示例讲解如何在 RoBERTa 预训练流程中以线性复杂度自注意力Wang et al., 2020替代标准 O(n²) 注意力。你将掌握 Linformer 的核心压缩原理、linformer_roberta模型注册方式、四个专有命令行参数的语义与默认值以及从数据处理、训练启动到模型加载的完整实操流程。背景为什么需要 Linformer标准 Transformer 自注意力的计算量与序列长度呈二次关系。对长度为 n 的序列注意力矩阵为 n×n其时间与空间复杂度均为 O(n²)。当 RoBERTa 这类模型被用于长文本预训练例如tokens_per_sample512乃至更长时注意力计算会成为显存与算力的主要瓶颈。LinformerWang et al., 2020论文见 arXiv:2006.04768的核心思想是通过两个可学习的线性投影矩阵把长度为 n 的 Key 和 Value 序列先压缩到固定长度 kk ≪ n再计算注意力。这样注意力矩阵的规模从 n×n 变为 n×k整体复杂度降为O(n)。fairseq 仓库的 examples/linformer/ 示例即是对该论文的官方实现可直接用于训练 Linformer 版的 RoBERTa 模型。一、快速上手三步训练 Linformer RoBERTa原文档指出训练 Linformer RoBERTa 与标准的 RoBERTa 预训练流程 基本一致唯一的差异在于训练命令需要额外指定--user-dir examples/linformer/linformer_src --arch linformer_roberta_base其中--user-dir告诉 fairseq 加载位于该目录下的自定义模型与模块注册代码--arch选择注册好的 Linformer 架构。下面按 RoBERTa 预训练的三步流程完整展开。1预处理数据数据需按语言模型格式组织每个文档之间用空行分隔配合--sample-break-mode complete_doc使用此时每行拼接成 1D 文本流。以 WikiText-103 为例wget https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-raw-v1.zip unzip wikitext-103-raw-v1.zip用 GPT-2 BPE 编码mkdir -p gpt2_bpe wget -O gpt2_bpe/encoder.json https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/encoder.json wget -O gpt2_bpe/vocab.bpe https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/vocab.bpe for SPLIT in train valid test; do \ python -m examples.roberta.multiprocessing_bpe_encoder \ --encoder-json gpt2_bpe/encoder.json \ --vocab-bpe gpt2_bpe/vocab.bpe \ --inputs wikitext-103-raw/wiki.${SPLIT}.raw \ --outputs wikitext-103-raw/wiki.${SPLIT}.bpe \ --keep-empty \ --workers 60; \ done然后二值化wget -O gpt2_bpe/dict.txt https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/dict.txt fairseq-preprocess \ --only-source \ --srcdict gpt2_bpe/dict.txt \ --trainpref wikitext-103-raw/wiki.train.bpe \ --validpref wikitext-103-raw/wiki.valid.bpe \ --testpref wikitext-103-raw/wiki.test.bpe \ --destdir>DATA_DIRdata-bin/wikitext-103 fairseq-hydra-train -m --config-dir examples/roberta/config/pretraining \ --config-name base task.data$DATA_DIR要切换到 Linformer只需在命令中加入--user-dir与--arch参数fairseq-hydra-train -m --config-dir examples/roberta/config/pretraining \ --config-name base task.data$DATA_DIR \ --user-dir examples/linformer/linformer_src \ --arch linformer_roberta_base也可以继续沿用 RoBERTa 的经典命令行训练入口将上述两个参数同样传给fairseq-train。关于资源与超参数原预训练文档给出如下提示上述命令假设在 8×32GB V100 GPU 上训练每张卡 batch size 为 16dataset.batch_size并以 16 倍梯度累积optimization.update_freq将总 batch size 扩到 2048 个序列GPU 更少或显存更小时应调小dataset.batch_size、调大update_freq补偿。学习率与 batch size 紧密相关增大 batch size 时通常需同步调大学习率具体取值随数据集而定batch sizepeak learning rate2560.000120480.000581920.0007在 base.yaml 中masked LM 预训练默认tokens_per_sample: 512、model.max_positions: 512、model.dropout / attention_dropout: 0.1、Adam polynomial_decay 调度、lr: [0.0005]、max_update: 125000。这些默认配置与 Linformer 的max_seq_len语义直接相关见下文源码分析。3加载预训练模型from fairseq.models.roberta import RobertaModel roberta RobertaModel.from_pretrained(checkpoints, checkpoint_best.pt, path/to/data) assert isinstance(roberta.model, torch.nn.Module)二、Linformer 架构的四把钥匙命令行参数详解在 linformer_roberta.py 中LinformerModel.add_args在继承 RoBERTa 全部参数的基础上新增了四个参数它们共同控制如何把长序列压缩成短序列参数类型默认值含义--compressedint4序列长度的压缩倍率压缩后长度compress_seq_len max_seq_len // compressed--shared-kv-compressedint0在同一层内 K 与 V 共享同一个压缩投影矩阵为 1 时共享--shared-layer-kv-compressedint0K 与 V 的压缩矩阵跨所有层共享为 1 时共享--freeze-compressint0冻结压缩层的参数不参与梯度更新为 1 时冻结这四个参数在base_architecture中被显式赋默认值linformer_roberta.pyargs.compressed getattr(args, compressed, 4) args.shared_kv_compressed getattr(args, shared_kv_compressed, 0) args.shared_layer_kv_compressed getattr(args, shared_layer_kv_compressed, 0) args.freeze_compress getattr(args, freeze_compress, 0) roberta_base_architecture(args)即默认以 4 倍压缩比、K/V 各自独立投影、不冻结压缩层的方式训练。在--shared-layer-kv-compressed 1时压缩矩阵由编码器统一创建并注入每一层见 linformer_sentence_encoder.py 的build_encoder_layer该矩阵以max_positions → max_positions // compressed的nn.Linear初始化并使用gain1/√2的 Xavier 初始化同样支持freeze_compress冻结。三、源码级实现从模型注册到线性注意力前向1模型注册与继承linformer_roberta.py 通过装饰器完成注册register_model(linformer_roberta) class LinformerModel(RobertaModel): ... classmethod def build_model(cls, args, task): base_architecture(args) if not safe_hasattr(args, max_positions): args.max_positions args.tokens_per_sample encoder LinformerEncoder(args, task.source_dictionary) return cls(args, encoder)LinformerModel直接继承RobertaModel因此预训练/加载 API 完全兼容LinformerEncoder继承RobertaEncoder其build_encoder返回LinformerTransformerEncoder并对参数执行init_bert_paramsBERT 风格初始化架构注册了三个别名linformer_roberta、linformer_roberta_base12 层、768 维、12 头、FFN 3072继承自 RoBERTa base见 fairseq/models/roberta/model.py 的base_architecture与linformer_roberta_large24 层、1024 维、16 头、FFN 4096继承自roberta_large_architecture。从源码结构可以推断要训练 Linformer Large 版只需将--arch换成linformer_roberta_large。2双向编码器与压缩层共享linformer_sentence_encoder.py 中的LinformerTransformerEncoder继承标准TransformerEncoder用于 BERT/XLM 风格的双向句子编码。其输出与 RoBERTa 一致内部各层状态列表每张 T×B×C与首 tokenCLS的句子表示B×C。当开启shared_layer_kv_compressed时跨层共享的压缩层在首个build_encoder_layer调用中创建后续所有层复用同一实例。3线性注意力前向压缩发生在哪一步核心实现在 multihead_linear_attention.py 的MultiheadLinearAttention.forward中。以 self-attention 分支为例计算顺序为# 1) Query 正常投影 q self.q_proj(query) # 2) Key 先做序列压缩再做投影 k_input query.permute(1, 2, 0).contiguous() # B * C * T k_input F.linear(k_input, self.compress_k.weight[:, 0:tgt_len]).permute(2, 0, 1).contiguous() k self.k_proj(k_input) # 3) Value 同理若 shared_kv_compressed 1 则复用 compress_k v_input query.permute(1, 2, 0).contiguous() v_input F.linear(v_input, self.compress_v.weight[:, 0:tgt_len]).permute(2, 0, 1).contiguous() v self.v_proj(v_input)要点拆解先压缩、后投影compress_k/compress_v是形状为(compress_seq_len, max_seq_len)的线性层把B×C×T的输入压缩为B×C×(T/compressed)随后 K/V 投影层在压缩后的短序列上工作。注意力矩阵随之从 n×n 降为 n×k实现 O(n) 复杂度动态裁剪weight[:, 0:tgt_len]压缩矩阵按max_seq_len来自args.max_positions初始化实际前向时只取前tgt_len列因此实际序列短于max_positions时也完全正确推理支持模块通过with_incremental_state装饰并实现了reorder_incremental_state与 key/value 缓存逻辑可配合增量解码encoder_decoder_attention分支K/V 由显式 key 提供不做压缩保留了对 encoder-decoder 注意力的兼容add_zero_attn、attn_mask、key_padding_mask、before_softmax、need_head_weights等接口与 fairseq 标准MultiheadAttention保持一致便于无缝替换。4编码器层与 checkpoint 兼容linformer_sentence_encoder_layer.py 的LinformerTransformerEncoderLayer继承标准TransformerEncoderLayer仅重写build_self_attention返回MultiheadLinearAttention并把压缩层的参数、头数、dropout、quant_noise 等参数逐项传入。该层与LinformerEncoder都实现了upgrade_state_dict_named用于加载旧版本 checkpoint旧 checkpoint 中若压缩层是跨层错误共享的shared_compress_layer.weight存在会自动重建不共享的压缩层以匹配旧行为LinformerEncoder还会检查嵌入层与 LM head 是否真正权重绑定若绑定错误则重建无权重共享的lm_head模块层面兼容旧的in_proj_weight/in_proj_bias格式将其拆分还原为独立的q_proj / k_proj / v_proj权重。这意味着从早期训练产出的 Linformer 权重仍可被当前代码正确加载。四、测试与验证线索虽然 tests/ 目录没有针对 Linformer 的独立测试文件但从源码结构可以推断其正确性由三方面保障接口对齐MultiheadLinearAttention.forward的输入输出签名与 fairseq 标准多头注意力完全一致need_weights / attn_mask / incremental_state等行为均有对应实现参数化注册linformer_roberta_base复用 RoBERTa 的全部超参数默认值任何对标准 RoBERTa 有效的训练配置均可直接叠加--user-dir与--arch使用数值路径自洽压缩矩阵在初始化、共享、冻结三种模式下分别有对应的初始化与requires_grad逻辑前向中 K/V 的压缩路径共享或独立与权重定义一一对应。建议在自有小数据集上先用--arch linformer_roberta_base跑通上述三步流程再逐步放大--compressed与序列长度验证显存与速度收益。五、引用若使用该工作请引用article{wang2020linformer, title{Linformer: Self-Attention with Linear Complexity}, author{Wang, Sinong and Li, Belinda and Khabsa, Madian and Fang, Han and Ma, Hao}, journal{arXiv preprint arXiv:2006.04768}, year{2020} }延伸阅读Linformer 示例目录模型与模块源码入口models/linformer_roberta.py、modules/multihead_linear_attention.pyRoBERTa 预训练教程数据处理、训练与加载的完整步骤RoBERTa 模型实现roberta_base_architecture等超参数默认值RoBERTa 预训练配置masked LM 训练默认超参数【免费下载链接】fairseqFacebook AI Research Sequence-to-Sequence Toolkit written in Python.项目地址: https://gitcode.com/gh_mirrors/fa/fairseq创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考