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/transformersYOSOYou 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 采样方案在 GLUE512 序列长度上取得与标准预训练 Transformer 相当的效果在 Long Range ArenaLRA长序列基准上取得与 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.FunctionYosoCumulation期望版本YOSO-E不真正采样而是用闭式期望逼近注意力。YosoLSHCumulation采样版本真正执行哈希与基于 LSH 的累积需要自定义 CUDA 内核。两种运行模式YOSO-E 与 CUDA Kernel 版文档 模型文档 的 Usage tips 明确了关键的使用前提YOSO 注意力算法通过自定义 CUDA 内核实现——即以 CUDA C 编写、可在 GPU 上多次并行执行的函数内核提供了fast_hash函数它借助Fast Hadamard Transform快速哈达玛变换来近似 query、key 的随机投影利用这些哈希码lsh_cumulation函数再通过基于 LSH 的伯努利采样来近似自注意力若要启用自定义内核需设置config.use_expectation False为确保内核成功编译必须安装版本匹配的 PyTorch 与 cudatoolkit默认情况下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, version1) lsh_cumulation yoso.lsh_cumulation而YosoSelfAttention.__init__modeling_yoso.py中只有在满足is_torch_cuda_available()且is_ninja_available()且内核尚未加载时才尝试加载失败时仅打印 warning 并退回相应回退路径。前端掩码与张量处理在YosoSelfAttention.forwardmodeling_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_expectationFalse且head_dim 32时会把 query/key/value 在最后一维用零补齐到 32modeling_yoso.py算完后再裁剪回原始head_dim期望模式下use_expectation为真或训练阶段会对 query/key 做 L2 归一化nn.functional.normalize(p2, dim-1)最终输出的 context 也会再一次归一化若配置了conv_window还会对 value 做一次深度可分离卷积depth-wise conv2dkernel 为(conv_window, 1)并把卷积结果残差式叠加到注意力输出上modeling_yoso.py用于注入局部性先验。期望模式 YOSO-E为什么可以“不求导采样”YosoCumulation.forwardmodeling_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.forwardmodeling_yoso.py走真正的采样路径若use_fast_hashTrue调用内核的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.pylsh_backwardTrue利用内核的lsh_weighted_cumulation计算 query/key/value 的梯度保持整条链路“线性”lsh_backwardFalse退化为期望形式YOSO-E 的反向公式做稠密梯度计算。YosoConfig 配置参数详解YosoConfig定义在 configuration_yoso.py其中 YOSO 专有参数语义如下参数代码默认值含义use_expectationTrue是否使用 YOSO-E期望模式。为True时不需要编译 CUDA 内核文档同时注明它会覆盖num_hash的效果。hash_code_len9哈希函数生成的哈希码长度位宽决定哈希表容量2^hash_code_len。num_hash64在 YosoSelfAttention 中使用的哈希函数数量。conv_windowNone深度可分离卷积的核宽为None表示不使用卷积局部增强。use_fast_hash源码字段默认True配置内 docstring 文字写的是False是否使用基于哈达玛变换快速随机投影的自定义 CUDA 内核。注意该参数只在use_expectationFalse的采样路径中起作用若编译/加载内核失败代码会回退到纯 PyTorch 的hashing()。lsh_backwardTrue反向传播是否也走 LSH内核的lsh_weighted_cumulation为False时用期望形式的稠密梯度。除 YOSO 专有参数外YosoConfig对齐 BERT 风格的 BERT-base-4096 结构默认值如下均可通过YosoConfig(...)覆盖词表vocab_size 50265隐藏层hidden_size 768编码器层数num_hidden_layers 12注意力头num_attention_heads 12FFN 中间维度intermediate_size 3072激活函数hidden_act geluhidden_dropout_prob attention_probs_dropout_prob 0.1max_position_embeddings 4096即yoso-4096的由来type_vocab_size 1layer_norm_eps 1e-12initializer_range 0.02特殊 tokenpad_token_id 1、bos_token_id 0、eos_token_id 2tie_word_embeddings TrueMLM 输出解码器与输入嵌入共享权重见_tied_weights_keys。一个值得注意的编码细节在YosoEmbeddingsmodeling_yoso.py中位置 id 统一偏移2同时位置嵌入表大小取max_position_embeddings 2——这与 0/1/2 被 bos/pad/eos 占用的 token 约定相呼应。在 Transformers 中使用 YOSO实例化配置与模型参考YosoConfig类内示例configuration_yoso.pyfrom 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_expectationTrue) # 希望使用自定义 CUDA 内核fast_hash lsh_cumulation # 需要 GPU、正确版本的 PyTorch 与 cudatoolkit且环境具备 kernels/ninja 支持 config YosoConfig(use_expectationFalse, use_fast_hashTrue) model YosoModel(config)需要注意的适用前提文档建议按仓库中的环境检查utils.is_torch_cuda_available/is_ninja_availableutils 中提供确认内核可编译条件采样路径只有在这些条件满足时才会尝试加载kernels-community/yoso内核use_expectationFalse且head_dim 32时张量会被零填充对齐到 warp 大小后再进内核若只在 CPU 上做原型验证建议保持默认use_expectationTrueYOSO-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 风格结构编码器 各类输出头搭建类用途输出头结构YosoModel基础编码器输出 last hidden state 及可选 hidden states / attentions词嵌入 位置嵌入 12 层 YosoLayerYosoForMaskedLM掩码语言建模预训练目标YosoOnlyMLMHead解码权重与输入嵌入共享YosoForSequenceClassification句级分类 / 回归如 GLUEYosoClassificationHead取序列首 token等价 [CLS]接两层线性头YosoForMultipleChoice选择题/多项选择pre_classifier ReLU 单输出线性层YosoForTokenClassification序列标注token 级dropout 全连接到num_labelsYosoForQuestionAnswering抽取式问答SQuAD 风格qa_outputs输出 start/end logitsstart/end 两个交叉熵损失取平均这些输出类的核心实现细节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 2MLM 头与嵌入层通过_tied_weights_keys声明权重绑定关系便于resize_token_embeddings等操作自动同步。相关任务指南与验证资源YOSO 可以直接套用本仓库面向下游任务的标准教程来微调与评估这些任务向导覆盖了数据加载、指标、训练循环的完整最佳实践文本分类任务指南Token 分类任务指南问答任务指南掩码语言建模任务指南多项选择任务指南自动化测试位于 tests/models/yoso/test_modeling_yoso.py覆盖了上述各任务头的模型等价性、输出形状与 loss 前向/反向等常规校验可以作为你自定义参数或接入新下游任务的回归参考。总结与使用建议YOSO 在本仓库中的实现可以概括为一条主线、两种模式一条主线把“n×n 稠密注意力”替换为“哈希 按桶累积”从原理上把复杂度降到线性对 GPU 的适配体现为面向 warp 的 32 维对齐、以哈达玛变换实现的高速随机投影内核以及可选的深度卷积局部增强两种模式默认的YOSO-Euse_expectationTrue以闭式期望替代真实采样纯 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考