企业官网建设流程全解析

Transformers 中的 YOSO 模型：基于 LSH 与伯努利采样的线性复杂度自注意力实现解析

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

YOSO（You Only Sample (Almost) Once）是一类把 softmax 自注意力的二次复杂度降为线性的高效注意力模型。它在本仓库中由论文作者贡献，并在src/transformers/models/yoso/下提供了完整的 PyTorch 实现、配置类与 CUDA 内核接入逻辑。阅读本文后，你将理解 YOSO 的伯努利采样 + 局部敏感哈希（LSH）原理、YosoConfig全部参数的语义，以及如何在本项目中选择“纯 PyTorch 的 YOSO-E 模式”或“自定义 CUDA 内核模式”进行推理与训练。

模型背景与论文出处

YOSO 模型由 Zhanpeng Zeng、Yunyang Xiong、Sathya N. Ravi、Shailesh Acharya、Glenn Fung、Vikas Singh 提出，论文题为You Only Sample (Almost) Once: Linear Cost Self-Attention Via Bernoulli Sampling。该模型于 2021-11-18 被收录进 Hugging Face papers，并于 2022-01-26 由贡献者 novice03 合入本仓库（见 模型文档）。

论文的核心观点是：

• Transformer 的中心是自注意力机制，它捕获输入序列中 token 两两之间的交互，计算代价随序列长度呈二次方增长，导致超长序列上的训练代价高昂；
• YOSO 通过基于 LSH 的伯努利采样注意力机制，把该二次复杂度降为线性；
• 其做法是把自注意力看作与每个 token 相关联的伯努利随机变量之和，这些随机变量“原则上可由一次哈希全部采样”（实际中可能是某个小的常数次）；
• 论文给出了面向 GPU 架构专门改造的 LSH 采样方案，在 GLUE（512 序列长度）上取得与标准预训练 Transformer 相当的效果，在 Long Range Arena（LRA）长序列基准上取得与 softmax 自注意力一致的精度，同时带来显著的加速与内存节省。

对应的模型类型标识为model_type = "yoso"，默认对齐的预训练结构是uw-madison/yoso-4096（对应max_position_embeddings = 4096）。

整体架构与核心机制

为什么需要“一次采样”：把注意力写成伯努利变量之和

标准的 softmax 自注意力需要对序列中每个 query 与全部 key 计算相似度并归一化，复杂度为 O(n²)。YOSO 的思想是：把注意力权重解释为某种伯努利随机变量的期望，于是前向计算不必显式构造稠密的 n×n 注意力矩阵，而只需通过 LSH 把“大概率相关的 query-key 对”聚到一起。

从本仓库的 modeling_yoso.py 源码看，这一机制被实现为两个自定义的torch.autograd.Function：

• YosoCumulation（期望版本，YOSO-E）：不真正采样，而是用闭式期望逼近注意力。
• YosoLSHCumulation（采样版本）：真正执行哈希与基于 LSH 的累积，需要自定义 CUDA 内核。

两种运行模式：YOSO-E 与 CUDA Kernel 版

文档 模型文档 的 Usage tips 明确了关键的使用前提：

1. YOSO 注意力算法通过自定义 CUDA 内核实现——即以 CUDA C++ 编写、可在 GPU 上多次并行执行的函数；
2. 内核提供了fast_hash函数，它借助Fast Hadamard Transform（快速哈达玛变换）来近似 query、key 的随机投影；利用这些哈希码，lsh_cumulation函数再通过基于 LSH 的伯努利采样来近似自注意力；
3. 若要启用自定义内核，需设置config.use_expectation = False；为确保内核成功编译，必须安装版本匹配的 PyTorch 与 cudatoolkit；
4. 默认情况下config.use_expectation = True，使用YOSO-E，此时不需要编译任何 CUDA 内核，因此也更便于在 CPU 等无 GPU 环境下直接跑通。

内核加载逻辑（源码级）

在 modeling_yoso.py 中，load_cuda_kernels()通过仓库的hub_kernels集成机制拉取并暴露内核：

def load_cuda_kernels(): global lsh_cumulation if not is_kernels_available(): raise ImportError("kernels is not installed, please install it with `pip install kernels`") from ...integrations.hub_kernels import get_kernel yoso = get_kernel("kernels-community/yoso", version=1) lsh_cumulation = yoso.lsh_cumulation

而YosoSelfAttention.__init__（modeling_yoso.py）中，只有在满足is_torch_cuda_available()且is_ninja_available()且内核尚未加载时才尝试加载，失败时仅打印 warning 并退回相应回退路径。

前端掩码与张量处理

在YosoSelfAttention.forward（modeling_yoso.py）中可以看到以下几个关键细节：

• query/key/value 由隐藏状态经三个nn.Linear投影，并切成num_attention_heads × attention_head_size的形状；
• 进入前向计算前，输入掩码先做1.0 + attention_mask / 10000.0的缩放，再沿注意力头维度重复并摊平成(batch_size * num_heads, seq_len)的整数掩码，供 LSH 累积内核当作 query/key mask 使用；
• GPU warp 对齐：源码注释指出“CUDA 内核在输入尺寸为 GPU warp 大小（32）的倍数时效率最高”。因此当use_expectation=False且head_dim < 32时，会把 query/key/value 在最后一维用零补齐到 32（modeling_yoso.py），算完后再裁剪回原始head_dim；
• 期望模式下（use_expectation为真）或训练阶段会对 query/key 做 L2 归一化（nn.functional.normalize(p=2, dim=-1)），最终输出的 context 也会再一次归一化；
• 若配置了conv_window，还会对 value 做一次深度可分离卷积（depth-wise conv2d，kernel 为(conv_window, 1)），并把卷积结果残差式叠加到注意力输出上（modeling_yoso.py），用于注入局部性先验。

期望模式 YOSO-E：为什么可以“不求导采样”

YosoCumulation.forward（modeling_yoso.py）给出了 YOSO-E 的核心数学形式。对归一化后的 query 与 key，它计算：

expectation = (1 - torch.acos(torch.matmul(query, key.transpose(-1, -2))) / math.pi) ** hash_code_len expectation = expectation * query_mask[:, :, None] * key_mask[:, None, :] cumulation_value = torch.matmul(expectation, value)

即用闭式公式(1 - arccos(q·k)/π)^hash_code_len刻画“两个向量在一次哈希中撞桶”的期望概率，再乘以掩码并用它加权 value。其反向传播（modeling_yoso.py）也是闭式可微的，因此整条 YOSO-E 路径完全不需要真正的随机采样，也不需要 CUDA 内核，天然支持任意设备上的训练与推理。

LSH 采样模式：真正的哈希路径

YosoLSHCumulation.forward（modeling_yoso.py）走真正的采样路径：

• 若use_fast_hash=True，调用内核的lsh_cumulation.fast_hash(...)，用快速哈达玛变换近似随机投影生成 query/key 哈希码；
• 否则退回纯 PyTorch 的hashing(query, key, num_hash, hash_code_len)函数：生成随机投影矩阵、投影、按符号二值化（投影 > 0 记为 1），并借助2 ** arange(hash_code_len)把每位拼成哈希桶编号，哈希表容量为2 ** hash_code_len；
• 随后调用lsh_cumulation.lsh_cumulation(...)，按哈希码把各位置归入桶内做累积，从而避免显式构造完整的 n×n 注意力矩阵。

反向传播由config.lsh_backward控制（modeling_yoso.py）：

• lsh_backward=True：利用内核的lsh_weighted_cumulation计算 query/key/value 的梯度，保持整条链路“线性”；
• lsh_backward=False：退化为期望形式（YOSO-E 的反向公式）做稠密梯度计算。

YosoConfig 配置参数详解

YosoConfig定义在 configuration_yoso.py，其中 YOSO 专有参数语义如下：

除 YOSO 专有参数外，YosoConfig对齐 BERT 风格的 BERT-base-4096 结构，默认值如下（均可通过YosoConfig(...)覆盖）：

• 词表vocab_size = 50265，隐藏层hidden_size = 768，编码器层数num_hidden_layers = 12，注意力头num_attention_heads = 12；
• FFN 中间维度intermediate_size = 3072，激活函数hidden_act = "gelu"；
• hidden_dropout_prob = attention_probs_dropout_prob = 0.1；
• max_position_embeddings = 4096（即yoso-4096的由来），type_vocab_size = 1；
• layer_norm_eps = 1e-12，initializer_range = 0.02；
• 特殊 token：pad_token_id = 1、bos_token_id = 0、eos_token_id = 2；
• tie_word_embeddings = True（MLM 输出解码器与输入嵌入共享权重，见_tied_weights_keys）。

一个值得注意的编码细节在YosoEmbeddings（modeling_yoso.py）中：位置 id 统一偏移+2，同时位置嵌入表大小取max_position_embeddings + 2——这与 0/1/2 被 bos/pad/eos 占用的 token 约定相呼应。

在 Transformers 中使用 YOSO

实例化配置与模型

参考YosoConfig类内示例（configuration_yoso.py）：

from transformers import YosoConfig, YosoModel # 初始化一个 uw-madison/yoso-4096 风格的配置 configuration = YosoConfig() # 用该配置初始化（随机权重）模型 model = YosoModel(configuration) # 访问模型配置 configuration = model.config

直接切换 YOSO-E / CUDA 内核模式

from transformers import YosoConfig, YosoModel # 默认即 YOSO-E：无需编译 CUDA 内核，CPU 上也可运行 config = YosoConfig(use_expectation=True) # 希望使用自定义 CUDA 内核（fast_hash + lsh_cumulation）： # 需要 GPU、正确版本的 PyTorch 与 cudatoolkit，且环境具备 kernels/ninja 支持 config = YosoConfig(use_expectation=False, use_fast_hash=True) model = YosoModel(config)

需要注意的适用前提：

• 文档建议按仓库中的环境检查utils.is_torch_cuda_available/is_ninja_available（utils 中提供）确认内核可编译条件；采样路径只有在这些条件满足时才会尝试加载kernels-community/yoso内核；
• use_expectation=False且head_dim < 32时，张量会被零填充对齐到 warp 大小后再进内核；
• 若只在 CPU 上做原型验证，建议保持默认use_expectation=True（YOSO-E），这是无需任何 GPU 依赖的路径。

从原始 YOSO checkpoint 转换

仓库还提供了把论文原版 PyTorch 权重转换为 Transformers 格式的脚本 convert_yoso_pytorch_to_pytorch.py，核心函数签名如下（可结合 checkpoint 的max_position_embeddings使用）：

• rename_key(orig_key)：把原始命名空间的键名映射到 Transformers 前缀体系；
• convert_checkpoint_helper(max_position_embeddings, orig_state_dict)：按层改写状态字典；
• convert_yoso_checkpoint(checkpoint_path, yoso_config_file, pytorch_dump_path)：加载原始权重与YosoConfig并导出标准模型文件。

模型家族与支持的子任务

src/transformers/models/yoso/内的实现覆盖了常见的“预训练 + 微调”全流程，全部基于 BERT 风格结构（编码器 + 各类输出头）搭建：

这些输出类的核心实现细节（loss 计算方式、masked token 处理、QA 越界 clamp 等）可在 modeling_yoso.py 中逐一对照。

其中几个值得留意的实现事实：

• YosoForSequenceClassification使用features[:, 0, :]取首个 token 作为句向量（见YosoClassificationHead），与 BERT 的[CLS]惯例一致；
• YosoForMultipleChoice需要输入形状(batch_size, num_choices, seq_len)，内部会摊平成(-1, seq_len)后喂给编码器；
• YosoForQuestionAnswering在构造时会强制config.num_labels = 2；
• MLM 头与嵌入层通过_tied_weights_keys声明权重绑定关系，便于resize_token_embeddings等操作自动同步。

相关任务指南与验证资源

YOSO 可以直接套用本仓库面向下游任务的标准教程来微调与评估（这些任务向导覆盖了数据加载、指标、训练循环的完整最佳实践）：

• 文本分类任务指南
• Token 分类任务指南
• 问答任务指南
• 掩码语言建模任务指南
• 多项选择任务指南

自动化测试位于 tests/models/yoso/test_modeling_yoso.py，覆盖了上述各任务头的模型等价性、输出形状与 loss 前向/反向等常规校验，可以作为你自定义参数或接入新下游任务的回归参考。

总结与使用建议

YOSO 在本仓库中的实现可以概括为一条主线、两种模式：

• 一条主线：把“n×n 稠密注意力”替换为“哈希 + 按桶累积”，从原理上把复杂度降到线性；对 GPU 的适配体现为面向 warp 的 32 维对齐、以哈达玛变换实现的高速随机投影内核，以及可选的深度卷积局部增强；
• 两种模式：默认的YOSO-E（use_expectation=True）以闭式期望替代真实采样，纯 PyTorch、可移植、无需编译；需要真实采样与最大加速时，把use_expectation设为False并保证 GPU 环境与版本匹配的 PyTorch/cudatoolkit（以及仓库所需的 kernels/ninja 支撑），即可走fast_hash + lsh_cumulation的完整 CUDA 内核路径。

对于想快速验证 YOSO 长序列收益的读者，建议优先从默认 YOSO-E 配置出发跑通小规模样例，再按需切换内核模式，并结合本文列出的配置项、源码位置与任务指南逐步深入。

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers