Kosmos-2 内嵌 fairseq 的 Adaptive Input 语言模型训练指南：基于 Baevski Auli (2018) 的 Transformer LM 实践-创锋一号

Kosmos-2 内嵌 fairseq 的 Adaptive Input 语言模型训练指南：基于 Baevski & Auli (2018) 的 Transformer LM 实践

【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm

导读

本文聚焦于 Kosmos-2 仓库内嵌的 fairseq 框架中，如何复现 Baevski & Auli (2018) 提出的 Adaptive Input Representations 语言模型。你将学会：理解 Adaptive Input（自适应输入嵌入）与 Adaptive Softmax（自适应输出层）这对"输入/输出双侧"分层的核心机制；在 WikiText-103 与 Google Billion Words 数据集上从零训练自适应 Transformer 语言模型；以及通过源码理解--criterion adaptive_loss、--arch transformer_lm_wiki103、cutoff 划分等关键参数的底层实现原理。

关联文档：kosmos-2/fairseq/examples/language_model/README.adaptive_inputs.md，本文以其为骨架，并结合仓库内 fairseq 源码逐层展开。

一、背景：为什么需要 Adaptive Input Representations

大规模语言模型的词汇表往往达到数十万量级。传统做法为每个 token 维护一个高维 Embedding（输入侧），并在输出侧对全词表计算 Softmax 归一化。其代价是：

参数量爆炸：vocab_size × hidden_dim的输入嵌入矩阵与输出投影矩阵占用了模型绝大部分参数；
计算浪费：高频词与低频词共享同一维度，无法按词频弹性分配容量。

Baevski & Auli (2018) 的方案是"分层分配"：把词表按频率切分成若干 band（频带），高频词使用更高维的嵌入空间，低频词使用更低维的嵌入空间，再通过线性投影统一到模型维度；输出侧使用对应的 Adaptive Softmax（Grave et al., 2017）进行分层归一化。这样既能显著降低参数量与计算量，又能保持甚至提升困惑度表现。

本文讨论的 fairseq 实现在仓库 kosmos-2/fairseq/ 中，所有路径以下文为准。

二、预训练模型速览

文档给出了两个官方发布的 Adaptive Input 预训练模型（原文表格完整继承）：

说明	参数量	数据集	模型与测试集
Adaptive Inputs（Baevski and Auli, 2018）	1026M	Google Billion Words	可下载（.tar.bz2 压缩包，含模型与测试集）
Adaptive Inputs（Baevski and Auli, 2018）	247M	WikiText-103	可下载（.tar.bz2 压缩包，含模型与测试集）

对应的预训练模型注册名可以在 language_model README 中看到：transformer_lm.gbw.adaptive_huge（1026M，GBW）与transformer_lm.wiki103.adaptive（247M，WikiText-103）。下载解压后，可通过 PyTorch Hub 或TransformerLanguageModel.from_pretrained加载（详见该 README 的 Example usage 一节）。

三、数据准备：WikiText-103 预处理全流程

Adaptive Inputs 文档明确指出：预处理流程与通用语言模型 README 完全一致，因此这里完整展开 README.md 中的步骤。

3.1 下载并解压数据集

仓库提供了现成脚本 prepare-wikitext-103.sh，其逻辑是：从 S3 下载wikitext-103-v1.zip，若文件已存在则跳过下载，随后根据扩展名自动解压（.zip用 unzip，.tgz/.tar用 tar）：

cd examples/language_model/ bash prepare-wikitext-103.sh cd ../..

3.2 二值化（binarize）数据

下载解压后会得到wiki.train.tokens、wiki.valid.tokens、wiki.test.tokens三个分词文件。接下来用fairseq-preprocess将其转为 fairseq 的二进制格式：

TEXT=examples/language_model/wikitext-103 fairseq-preprocess \ --only-source \ --trainpref $TEXT/wiki.train.tokens \ --validpref $TEXT/wiki.valid.tokens \ --testpref $TEXT/wiki.test.tokens \ --destdir>fairseq-train --task language_modeling \ >if cfg.ddp_backend in {"c10d", "pytorch_ddp"}: raise Exception( "AdaptiveLoss is not compatible with the PyTorch " "version of DistributedDataParallel. Please use " "`--ddp-backend=legacy_ddp` instead." )

即：只要使用--criterion adaptive_loss，就不能使用 PyTorch 的 DDP 后端，必须显式传--ddp-backend=legacy_ddp，否则训练直接抛异常。

4.3 显存不足时的通用调优手法

沿用 README.md 的官方建议：显存溢出时降低--max-tokens（每 batch token 数）或--tokens-per-sample（序列长度），或调大--update-freq用梯度累积模拟更多 GPU 的效果。

五、架构级解读：`transformer_lm_wiki103`与 Adaptive Input 的源码实现

5.1 架构配置文件

--arch transformer_lm_wiki103对应的实际配置在 transformer_lm_baevski_wiki103.yaml：

decoder_embed_dim: 1024 decoder_output_dim: 1024 decoder_input_dim: 1024 decoder_ffn_embed_dim: 4096 decoder_layers: 16 decoder_attention_heads: 8 decoder_normalize_before: true no_decoder_final_norm: true adaptive_softmax_cutoff: "20000,60000" adaptive_softmax_dropout: 0.2 adaptive_softmax_factor: 4 adaptive_input: true adaptive_input_factor: 4 adaptive_input_cutoff: "20000,60000" tie_adaptive_weights: true tie_adaptive_proj: true

关键参数与含义：

参数	取值	含义
`decoder_embed_dim`	1024	解码器隐藏维度
`decoder_layers`	16	Transformer 解码器层数
`decoder_attention_heads`	8	注意力头数
`adaptive_input: true`	—	开启自适应输入嵌入
`adaptive_input_cutoff: "20000,60000"`	—	输入侧 cutoff 切分点
`adaptive_softmax_cutoff: "20000,60000"`	—	输出侧 cutoff 切分点（与输入一致以便权重共享）
`adaptive_softmax_factor: 4`	4	每往后一个 band，维度除以 4
`adaptive_softmax_dropout`	0.2	分层输出层间的 dropout
`tie_adaptive_weights: true`	—	输入输出各 band 权重共享
`tie_adaptive_proj: true`	—	各 band 的投影矩阵也共享（转置复用）

而 GBW 的 transformer_lm_baevski_gbw.yaml 则是decoder_embed_dim: 512、12 层、16 头，且adaptive_input: false、tie_adaptive_weights: false，体现了两套预训练配置的差异。

5.2 权重共享的硬约束（源码断言）

在 transformer_lm.py 中，当开启tie_adaptive_weights时有一组必须同时满足的断言：

if args.tie_adaptive_weights: assert args.adaptive_input assert args.adaptive_input_factor == args.adaptive_softmax_factor assert args.adaptive_softmax_cutoff == args.adaptive_input_cutoff assert args.decoder_input_dim == args.decoder_output_dim

这意味着：一旦启用权重共享，输入嵌入的 factor、cutoff 必须与输出侧完全一致，输入输出维度也必须相同，否则训练会直接断言失败。这正是"输入侧 AdaptiveInput 与输出侧 AdaptiveSoftmax 必须成对设计"的实现证据。

5.3 输入侧：AdaptiveInput 的分层嵌入

实现在 adaptive_input.py，核心逻辑：

cutoff 自动补全：若vocab_size > cutoff[-1]，自动在 cutoff 末尾追加vocab_size（L28-L33），因此传入"20000,60000"后，WikiText-103 的约 26 万词表会变成[20000, 60000, vocab_size]三段；
分层嵌入：对第 i 个 band，嵌入维度为int(initial_dim // factor ** i)（L43），即高频带用满维initial_dim，低频带逐级除以 factor（默认 4）；
统一投影：每个 band 的nn.Embedding之后接一个无 bias 的nn.Linear(dim, output_dim)（L44-L49），把各 band 的低维向量统一投影回模型维度output_dim；
前向分派：forward中按 token 值区间把输入拆到不同 band，各自过嵌入+投影后再写回结果张量（L69-L80）。

此外，init_weights对每个 band 的 Embedding 使用均值为 0、标准差为dim^-0.5的正态初始化，并把 padding 位置权重置 0；Linear 投影层使用 xavier_uniform 初始化。

5.4 输出侧：AdaptiveSoftmax 的分层打分

实现在 adaptive_softmax.py，要点：

head + tail 结构：head负责前cutoff[0]个高频词与若干"簇类"标签；tail是若干个子网络，每个负责一个低频 band（L129-L173）；
adapt_target 标签改写：训练时并不会对所有词打分，而是把低频词的目标改写为"簇类标签"，仅对命中的 band 计算局部 softmax（L180-L203），这是自适应 softmax 省算力的关键；
权重绑定：当传入adaptive_inputs时，TiedHeadModule与各 tail 通过TiedLinear直接复用输入侧嵌入权重（TiedLinear可选择是否转置，对应tie_proj），从而在源码层面印证了tie_adaptive_weights/tie_adaptive_proj的语义（L100-L115）。

5.5 损失侧：adaptive_loss 的分段交叉熵

实现在 adaptive_loss.py，forward中调用adaptive_softmax(net_output[0], orig_target)得到分段 logits 与改写后的 targets，然后对每个非空段分别计算带ignore_index=padding_idx的交叉熵并求和（L71-L79）。reduce_metrics中按sample_size（默认取 token 数）归一化输出 loss，并派生 ppl 指标（L92-L114）。

六、训练完成后的评估

训练产出在--save-dir checkpoints/transformer_wikitext-103下，最佳检查点为checkpoint_best.pt。沿用通用 LM README 的评估命令：

fairseq-eval-lm>fairseq-train --task language_modeling \ >@inproceedings{ baevski2018adaptive, title={Adaptive Input Representations for Neural Language Modeling}, author={Alexei Baevski and Michael Auli}, booktitle={International Conference on Learning Representations}, year={2019}, url={https://openreview.net/forum?id=ByxZX20qFQ}, }

（说明：文末仅保留 bibtex 引用字段，原文档中的外部下载与论文外链因仓库内无法验证且不属于仓库证据，不再以超链接形式输出；预训练模型与数据集名称、参数规模均以文档原表为准。）

企业官网建设流程全解析