Ludwig 中的 entmax从 sparsemax 到可学习稀疏性的 α-entmax 概率映射与损失函数实战指南【免费下载链接】ludwigLow-code framework for building custom LLMs, neural networks, and other AI models项目地址: https://gitcode.com/gh_mirrors/lu/ludwig导读本文围绕 Ludwig 仓库内置的 entmax 实现 展开系统讲解一类稀疏概率映射及其配套损失函数entmax 家族以 Tsallis α-熵推广 softmax在保留可微性的同时让输出概率产生精确的零值从而在注意力、多标签分类、TabNet 特征选择等场景中天然具备稀疏性与可解释性。读完本文你将掌握sparsemax、entmax15、通用α-entmax含对 α 的可学习梯度的数学原理、PyTorch API 用法、损失函数设计以及它们如何在 Ludwig 的损失模块与 TabNet 组合器中落地配置。一、entmax 是什么稀疏版的 softmax 概率映射传统 softmax 将任意实数向量映射为严格正的、处处非零的概率分布而 entmax 是一族稀疏概率映射sparse probability mapping它把输入向量映射到单纯形simplex上并允许部分分量精确等于 0。这种稀疏性对注意力机制、多标签分类、序列到序列模型尤其有价值模型可以显式地忽略无关信息而不是分配给它们一个微小但非零的权重。正如 entmax README 所概括该包提供精确的 partial-sort 算法针对 1.5-entmax 与 2-entmax即 sparsemax两种特例基于二分法bisection的通用算法支持任意 α-entmax对 α 的梯度实现自适应、可学习的稀疏度控制。在数学上α-entmax 求解如下带约束的优化问题见 root_finding.py 源码注释max_p x, p - H_α(p) s.t. p ≥ 0, Σp 1其中H_α(p)是 Tsallis α-熵。α 的不同取值构成一条连续谱系α 取值对应映射稀疏程度α → 1softmax无稀疏性处处非零α 1.51.5-entmaxentmax15中等稀疏介于两者之间α 22-entmax即sparsemax最大程度稀疏等价于欧氏投影到单纯形sparsemax 的几何意义非常直观它是把向量 x 投影到概率单纯形上的欧氏投影min_p ||x - p||₂投影结果天然具有精确零分量。这一思想源自 Martins Astudillo2016随后由 Peters 等人2019推广为 1.5-entmax 与通用 α-entmax并用于稀疏序列到序列模型。二、安装与导入仓库内嵌的 entmax 包entmax 在 Ludwig 中以内部子包的形式存在于ludwig/utils/entmax/目录共五个文件init.py版本号与全部公开 API 的导出activations.pysparsemax、entmax15 的函数式实现与nn.Module封装root_finding.py通用 α-entmax 的二分法实现含对 α 的反向传播losses.py对应的稀疏损失函数与损失类README.md 与 LICENSEMIT 协议版权归属 DeepSPIN2019。独立使用时可以通过 PyPI 安装原版包pip install entmax而在 Ludwig 项目中只需从内置路径导入无需额外安装例如from ludwig.utils.entmax import sparsemax, entmax15, entmax_bisect from ludwig.utils.entmax import Sparsemax, Entmax15, EntmaxBisect from ludwig.utils.entmax.losses import ( sparsemax_loss, entmax15_loss, entmax_bisect_loss, SparsemaxLoss, Entmax15Loss, EntmaxBisectLoss, )完整 API 清单见init.py 的__all__涵盖函数式接口小写命名、损失函数与对应的nn.Module类大写命名。运行环境要求为 Python 3、PyTorch ≥ 1.0单元测试需要 pytest。三、快速上手sparsemax 与 entmax15 的直观对比entmax README 给出了一个非常直观的三路对比示例这里完整复现并补充解读。对输入向量x [-2, 0, 0.5]分别施加 softmax、sparsemax 与 entmax15import torch from torch.nn.functional import softmax from ludwig.utils.entmax import sparsemax, entmax15 x torch.tensor([-2, 0, 0.5]) print(softmax(x, dim0)) # tensor([0.0486, 0.3592, 0.5922]) print(sparsemax(x, dim0)) # tensor([0.0000, 0.2500, 0.7500]) print(entmax15(x, dim0)) # tensor([0.0000, 0.3260, 0.6740])观察输出softmax给三个位置都分配了非零概率0.0486 / 0.3592 / 0.5922即使-2明显是无关项sparsemaxα2把-2对应位置的概率精确置为 0其余两项按 0.25 / 0.75 分配entmax15α1.5同样把-2置零但概率分配更平滑0.326 / 0.674稀疏性介于 softmax 与 sparsemax 之间。从 activations.py 源码 可以看到两者的计算内核_sparsemax_threshold_and_support通过累计和与阈值公式计算最优阈值τ和支持集大小非零分量个数输出clamp(X - τ, min0)_entmax_threshold_and_support基于均值、均方与方差公式ss rho * (mean_sq - mean**2)delta (1 - ss) / rho计算 1.5-entmax 的阈值输出clamp(X - τ, min0) ** 2因此 entmax15 的输出是阈值的平方。两者的正向传播都采用了与 softmax 相同的数值稳定性技巧——先减去每行最大值X X - max_val避免大数值溢出。3.1 函数式接口与 nn.Module 封装sparsemax与entmax15同时提供函数式接口与模块化封装见 activations.pyfrom ludwig.utils.entmax import Sparsemax, Entmax15 # 函数式直接作用于张量 p sparsemax(x, dim-1) # 模块化便于嵌入 nn.Sequential 或自定义网络层 layer Sparsemax(dim-1, kNone) # kNone 表示全排序 p layer(x) layer Entmax15(dim-1, kNone) p layer(x)两者的nn.Module版本Sparsemax、Entmax15都接受dim沿哪个维度归一化与kpartial-sort 的元素数量详见第五节两个构造参数并在forward中透传self.training以保证推理阶段走 torchscript 兼容的快速路径详见第六节。四、通用 α-entmax二分法实现与可学习稀疏度4.1 支持任意 α 的二分求根算法当 α 不是 1.5 或 2 这样的特例时无法使用解析的 partial-sort 公式root_finding.py提供了基于**二分法bisection即二分搜索求根**的通用算法entmax_bisect。其核心思想是对每个向量在[τ_lo, τ_hi]区间内反复二分寻找使Σ p(x_i - τ) 1成立的阈值τ迭代n_iter次后得到近似投影最后通过ensure_sum_one归一化保证输出严格满足概率和为 1。函数签名与关键参数见 root_finding.pyentmax_bisect(X, alpha1.5, dim-1, n_iter50, ensure_sum_oneTrue, trainingTrue) sparsemax_bisect(X, dim-1, n_iter50, ensure_sum_oneTrue, trainingTrue)alphaα 取值1。alpha2等价于 sparsemax理论上alpha1会退化为 softmax但数值上不支持 α1需要 softmax 时请直接使用torch.nn.functional.softmax源码 docstring 对此有明确警告n_iter二分迭代次数float32 精度下约24 次即可达到机器精度默认 50 次留足余量ensure_sum_one是否对结果做归一化保证输出精确满足Σp 1防止下游数值问题alpha还支持传入torch.Tensor可以为每个样本行指定不同的 α形状需满足沿dim维度为 1、其余维度与 X 一致或可广播。此外还提供了 sparsemax 的二分特例SparsemaxBisectFunction通过重写_gpx ** (α-1)与_gp_invy ** (1/(α-1))为恒等映射省去幂运算开销见 root_finding.py 第 106-141 行。4.2 对 α 求梯度自适应稀疏度entmax_bisect是同时关于 X 和 α 可微的。EntmaxBisectFunction.backward在ctx.needs_input_grad[1]为真时利用香农熵项S Y * log(Y)与偏斜分布Y_skewed推导出d_alpha的闭式表达式见 root_finding.py 第 48-65 行注释特别提醒α 不能接近 1否则存在不确定性indetermination。这意味着我们可以把 α 当作可学习参数让模型自己决定注意力应该多稀疏。README 给出的对 α 求梯度的完整示例import torch from torch.autograd import grad from ludwig.utils.entmax import entmax_bisect x torch.tensor([[-1, 0, 0.5], [1, 2, 3.5]]) alpha torch.tensor(1.33, requires_gradTrue) p entmax_bisect(x, alpha) print(p) # tensor([[0.0460, 0.3276, 0.6264], # [0.0026, 0.1012, 0.8963]], grad_fnEntmaxBisectFunctionBackward) print(grad(p[0, 0], alpha)) # (tensor(-0.2562),)输出张量的grad_fn为EntmaxBisectFunctionBackward且p[0,0]对alpha的梯度可正常回传——这正是 Gradients w.r.t. alpha for adaptive, learned sparsity 这一特性的落地方式也是后续 TabNet 中adaptive模式的数学基础。五、深入实现partial-sort 算法与 k 参数对于 1.5-entmax 与 sparsemax 两个特例activations.py 采用精确的 partial-sort部分排序算法避免对整行做全量排序从而获得更优的常数因子性能。_sparsemax_threshold_and_support与_entmax_threshold_and_support都接受一个k参数kNone或k ≥ X.shape[dim]直接对全行排序torch.sort(..., descendingTrue)k X.shape[dim]改用torch.topk(X, kk, dimdim)只取最大的 k 个元素期望稀疏解的非零个数不超过 k递归兜底如果解的实际稀疏度超过 ksupport_size k说明支持集被截断算法会以2*k的倍增计划递归调用自身直到求出完整的支持集见_roll_last 递归分支。在 README 与 activations.py 的 docstring 中都强调k的取值应略大于预期的非零个数以获得最优性能。k同时作为函数式接口与Sparsemax/Entmax15模块的构造参数对外暴露。5.1 自定义反向传播的推导要点SparsemaxFunction.backward与Entmax15Function.backward都只依赖正向保存的输出张量即可完成梯度计算无需保存中间排序结果sparsemax对输出为 0 的位置梯度置零再对支持集内做去均值投影grad_input - v_hat这正是单纯形投影的雅可比形式entmax15利用gppr Y.sqrt()即1 / g(Y)计算加权修正实现dX - q * gppr的投影。两者都返回dX, None, None即只对输入 X 求梯度。六、训练与推理分离torchscript 兼容设计sparsemax、entmax15、entmax_bisect、sparsemax_bisect四个函数都带有trainingTrue关键字参数源码注释明确说明了原因自定义torch.autograd.Function无法被 torchscript 序列化activations.py 第 229-234 行引用了 PyTorch issue #22329。因此训练阶段trainingTrue走CustomFunction.apply(...)路径获得自定义反向传播与对 α 的梯度推理/导出阶段trainingFalse直接调用纯 PyTorch 算子组合的正向函数_sparsemax_forward/_entmax15_forward/_entmax_bisect_forward保证模型可被 torchscript 追踪导出。对应的nn.Module封装Sparsemax、Entmax15、EntmaxBisect、SparsemaxBisect在forward中传入self.training因此用户在模型eval()/train()切换时会自动选择正确路径无需手动干预。七、稀疏损失函数Fenchel-Young 损失家族entmax 不仅提供了稀疏概率映射还配套了对应的稀疏损失函数——它们是 softmax 交叉熵的稀疏替代品通过 Fenchel 对偶推导而来天然与对应激活函数匹配。全部实现位于 losses.pyREADME 与 loss_modules.py 的注释均将其称为激活函数 Ω 函数的 Fenchel 共轭损失。函数式接口一览函数对应映射特点sparsemax_loss(X, target, kNone)sparsemaxpartial-sort 计算ω (1 - Σp²) / 2sparsemax_bisect_loss(X, target, n_iter50)sparsemax二分法计算entmax15_loss(X, target, kNone)1.5-entmaxpartial-sort 计算ω (1 - Σ p·√p) / 0.75entmax_bisect_loss(X, target, alpha1.5, n_iter50)α-entmax二分法计算ω (1 - Σp^α) / (α(α-1))支持逐行 α它们的输入约定与交叉熵一致X为形状(n_samples, n_classes)的分数logitstarget为形状(n_samples,)的长整型标签满足0 ≤ target n_classes。_GenericLossFunction的正向传播非常优雅先用对应激活函数求出最优概率p_star计算正则项ω(p_star, α)再通过scatter_add_把目标类位置的贡献扣除并加上线性项einsum(ij,ij-i, p_star, X)最终得到逐样本损失反向传播只需保存p_star即可见 losses.py 第 39-64 行。除了函数式接口还提供了带ignore_index与reduction的nn.Module封装类SparsemaxLoss、SparsemaxBisectLoss、Entmax15Loss、EntmaxBisectLoss。它们继承自_GenericLossignore_index默认值为 Ludwig 的IGNORE_INDEX_TOKEN_ID可屏蔽填充位置如序列任务中的 pad token被屏蔽位置的损失被masked_fill_置零reduction支持elementwise_mean默认等价于求均值、sum、none三种模式非法值会抛出带提示的ValueError当所有行都被 ignore 时返回零损失与零梯度避免除零见 losses.py 第 31-34 行。from ludwig.utils.entmax.losses import Entmax15Loss criterion Entmax15Loss(k100) # 默认 k100可配合 ignore_index 使用 loss criterion(logits, targets) # logits: [B, C], targets: [B]八、在 Ludwig 中的实际集成损失模块与 TabNet8.1 注册为输出特征的损失函数Ludwig 的损失模块 loss_modules.py 将SparsemaxLoss与Entmax15Loss注册为可用损失类型通过register_loss装饰器对应SparsemaxLossConfig与Entmax15LossConfig。其 docstring 明确阐述了定位SparsemaxLosssoftmax 交叉熵的稀疏替代品……对 sparsemax 支持集之外的类别分配零梯度在概率单纯形上产生精确稀疏性参考文献为 Martins Astudillo (ICML 2016)Entmax15Lossα1.5 entmax 激活的 Fenchel 共轭损失……产生介于 softmax稠密与 sparsemax最大稀疏之间的中等稀疏概率分布参考文献为 Peters et al. (ACL 2019)。这两个损失类都实现了LogitsInputsMixin直接接收[B, C]的 logits 与[B]的长整型标签内部委托给ludwig.utils.entmax.losses中的底层实现。这意味着你可以在 Ludwig 配置中为分类输出特征直接指定type: sparsemax或type: entmax15作为损失函数无需编写任何自定义代码。8.2 TabNet 组合器中的 entmax 稀疏注意力entmax 在 Ludwig 中最引人注目的应用是TabNet 组合器的注意力模块。tabnet_modules.py 从ludwig.utils.entmax导入了Entmax15、EntmaxBisect、Sparsemax在AttentiveTransformer中依据entmax_mode选择稀疏映射if self.entmax_mode sparsemax: self.entmax_module Sparsemax() elif self.entmax_mode entmax15: self.entmax_module Entmax15() else: # adaptive self.entmax_module EntmaxBisect(alphaself.trainable_alpha)其中adaptive模式把 α 注册为可学习缓冲区self.register_buffer(trainable_alpha, torch.tensor(entmax_alpha, requires_gradTrue))正是利用了本文第四节介绍的对 α 求梯度能力让模型在训练中自动调节特征选择的稀疏程度。对应的配置 schema 位于 tabnet.py两个参数的约束如下entmax_modeStringOptions可选值[entmax15, sparsemax, constant, adaptive]默认sparsemax。文档说明entmax 是一族推广 softmax 与 sparsemax 的稀疏概率映射entmax_mode控制稀疏度entmax_alphaFloatRange取值范围1.0 到 2.0默认 1.5。schema 注释明确1对应 softmax2对应 sparsemax当entmax_mode为adaptive时它作为可学习参数的初始值。在 TabNet 配置文件中即可直接使用combiner: type: tabnet entmax_mode: adaptive entmax_alpha: 1.5九、测试验证数值一致性与梯度正确性仓库为 entmax 提供了完善的单元测试位于 tests/ludwig/utils/entmax/test_root_finding.py验证partial-sort 实现与二分法实现输出一致如test_sparsemax断言sparsemax(x, 1)与sparsemax_bisect(x)的平方误差小于1e-7test_entmax15同理并在训练/推理两种模式下都做了验证同时使用torch.autograd.gradcheck对sparsemax_bisect做梯度数值校验test_sparsemax_grad确认自定义反向传播的正确性test_losses.py验证各稀疏损失函数的数值行为test_topk.py 与 test_mask.py分别验证 partial-sort 的k参数路径与ignore_index掩码逻辑。此外test_tabnet_modules.py 覆盖了 TabNet 组合器中 entmax 相关参数的组合。这些测试既保证了数学实现的正确性也保证了集成进 Ludwig 后各配置组合的行为稳定。十、论文出处与进一步阅读entmax 的理论根基来自两篇核心论文原文以 BibTeX 形式记录于 READMEPeters, Niculae, and Martins, Sparse Sequence-to-Sequence Models, Proc. ACL 2019—— 提出 1.5-entmax 与通用 α-entmax并将其应用于序列到序列模型的注意力稀疏化Correia, Niculae, and Martins, Adaptively Sparse Transformers, Proc. EMNLP-IJCNLP 2019—— 提出对 α 的梯度计算实现自适应稀疏 Transformer。如需进一步理解README 还推荐了三篇延伸文献Blondel, Martins, and Niculae, 2019, Learning with Fenchel-Young Losses——本仓库损失函数的理论框架来源Martins and Astudillo, 2016, From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification——sparsemax 的原始出处Peters and Martins, 2019, IT-IST at the SIGMORPHON 2019 Shared Task: Sparse Two-headed Models for Inflection——entmax 在词形变化任务上的工程应用。结语entmax 家族为深度学习提供了一条从稠密到稀疏的连续调节路径sparsemax 给出最激进的精确稀疏entmax15 提供折中的半稀疏而基于二分法的 α-entmax 配合对 α 的可微性让稀疏度本身成为可学习的模型参数。在 Ludwig 中这套实现不仅以独立子包形式完整保留ludwig/utils/entmax/还深度集成了输出特征的稀疏损失sparsemax/entmax15与 TabNet 组合器的稀疏特征选择entmax_mode: sparsemax | entmax15 | adaptive。无论你是想在自定义模型中引入稀疏注意力还是在 Ludwig 配置中直接启用稀疏损失本文涉及的 API、参数与源码路径都可作为直接的实践依据。【免费下载链接】ludwigLow-code framework for building custom LLMs, neural networks, and other AI models项目地址: https://gitcode.com/gh_mirrors/lu/ludwig创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Ludwig 中的 entmax：从 sparsemax 到可学习稀疏性的 α-entmax 概率映射与损失函数实战指南
Ludwig 中的 entmax从 sparsemax 到可学习稀疏性的 α-entmax 概率映射与损失函数实战指南【免费下载链接】ludwigLow-code framework for building custom LLMs, neural networks, and other AI models项目地址: https://gitcode.com/gh_mirrors/lu/ludwig导读本文围绕 Ludwig 仓库内置的 entmax 实现 展开系统讲解一类稀疏概率映射及其配套损失函数entmax 家族以 Tsallis α-熵推广 softmax在保留可微性的同时让输出概率产生精确的零值从而在注意力、多标签分类、TabNet 特征选择等场景中天然具备稀疏性与可解释性。读完本文你将掌握sparsemax、entmax15、通用α-entmax含对 α 的可学习梯度的数学原理、PyTorch API 用法、损失函数设计以及它们如何在 Ludwig 的损失模块与 TabNet 组合器中落地配置。一、entmax 是什么稀疏版的 softmax 概率映射传统 softmax 将任意实数向量映射为严格正的、处处非零的概率分布而 entmax 是一族稀疏概率映射sparse probability mapping它把输入向量映射到单纯形simplex上并允许部分分量精确等于 0。这种稀疏性对注意力机制、多标签分类、序列到序列模型尤其有价值模型可以显式地忽略无关信息而不是分配给它们一个微小但非零的权重。正如 entmax README 所概括该包提供精确的 partial-sort 算法针对 1.5-entmax 与 2-entmax即 sparsemax两种特例基于二分法bisection的通用算法支持任意 α-entmax对 α 的梯度实现自适应、可学习的稀疏度控制。在数学上α-entmax 求解如下带约束的优化问题见 root_finding.py 源码注释max_p x, p - H_α(p) s.t. p ≥ 0, Σp 1其中H_α(p)是 Tsallis α-熵。α 的不同取值构成一条连续谱系α 取值对应映射稀疏程度α → 1softmax无稀疏性处处非零α 1.51.5-entmaxentmax15中等稀疏介于两者之间α 22-entmax即sparsemax最大程度稀疏等价于欧氏投影到单纯形sparsemax 的几何意义非常直观它是把向量 x 投影到概率单纯形上的欧氏投影min_p ||x - p||₂投影结果天然具有精确零分量。这一思想源自 Martins Astudillo2016随后由 Peters 等人2019推广为 1.5-entmax 与通用 α-entmax并用于稀疏序列到序列模型。二、安装与导入仓库内嵌的 entmax 包entmax 在 Ludwig 中以内部子包的形式存在于ludwig/utils/entmax/目录共五个文件init.py版本号与全部公开 API 的导出activations.pysparsemax、entmax15 的函数式实现与nn.Module封装root_finding.py通用 α-entmax 的二分法实现含对 α 的反向传播losses.py对应的稀疏损失函数与损失类README.md 与 LICENSEMIT 协议版权归属 DeepSPIN2019。独立使用时可以通过 PyPI 安装原版包pip install entmax而在 Ludwig 项目中只需从内置路径导入无需额外安装例如from ludwig.utils.entmax import sparsemax, entmax15, entmax_bisect from ludwig.utils.entmax import Sparsemax, Entmax15, EntmaxBisect from ludwig.utils.entmax.losses import ( sparsemax_loss, entmax15_loss, entmax_bisect_loss, SparsemaxLoss, Entmax15Loss, EntmaxBisectLoss, )完整 API 清单见init.py 的__all__涵盖函数式接口小写命名、损失函数与对应的nn.Module类大写命名。运行环境要求为 Python 3、PyTorch ≥ 1.0单元测试需要 pytest。三、快速上手sparsemax 与 entmax15 的直观对比entmax README 给出了一个非常直观的三路对比示例这里完整复现并补充解读。对输入向量x [-2, 0, 0.5]分别施加 softmax、sparsemax 与 entmax15import torch from torch.nn.functional import softmax from ludwig.utils.entmax import sparsemax, entmax15 x torch.tensor([-2, 0, 0.5]) print(softmax(x, dim0)) # tensor([0.0486, 0.3592, 0.5922]) print(sparsemax(x, dim0)) # tensor([0.0000, 0.2500, 0.7500]) print(entmax15(x, dim0)) # tensor([0.0000, 0.3260, 0.6740])观察输出softmax给三个位置都分配了非零概率0.0486 / 0.3592 / 0.5922即使-2明显是无关项sparsemaxα2把-2对应位置的概率精确置为 0其余两项按 0.25 / 0.75 分配entmax15α1.5同样把-2置零但概率分配更平滑0.326 / 0.674稀疏性介于 softmax 与 sparsemax 之间。从 activations.py 源码 可以看到两者的计算内核_sparsemax_threshold_and_support通过累计和与阈值公式计算最优阈值τ和支持集大小非零分量个数输出clamp(X - τ, min0)_entmax_threshold_and_support基于均值、均方与方差公式ss rho * (mean_sq - mean**2)delta (1 - ss) / rho计算 1.5-entmax 的阈值输出clamp(X - τ, min0) ** 2因此 entmax15 的输出是阈值的平方。两者的正向传播都采用了与 softmax 相同的数值稳定性技巧——先减去每行最大值X X - max_val避免大数值溢出。3.1 函数式接口与 nn.Module 封装sparsemax与entmax15同时提供函数式接口与模块化封装见 activations.pyfrom ludwig.utils.entmax import Sparsemax, Entmax15 # 函数式直接作用于张量 p sparsemax(x, dim-1) # 模块化便于嵌入 nn.Sequential 或自定义网络层 layer Sparsemax(dim-1, kNone) # kNone 表示全排序 p layer(x) layer Entmax15(dim-1, kNone) p layer(x)两者的nn.Module版本Sparsemax、Entmax15都接受dim沿哪个维度归一化与kpartial-sort 的元素数量详见第五节两个构造参数并在forward中透传self.training以保证推理阶段走 torchscript 兼容的快速路径详见第六节。四、通用 α-entmax二分法实现与可学习稀疏度4.1 支持任意 α 的二分求根算法当 α 不是 1.5 或 2 这样的特例时无法使用解析的 partial-sort 公式root_finding.py提供了基于**二分法bisection即二分搜索求根**的通用算法entmax_bisect。其核心思想是对每个向量在[τ_lo, τ_hi]区间内反复二分寻找使Σ p(x_i - τ) 1成立的阈值τ迭代n_iter次后得到近似投影最后通过ensure_sum_one归一化保证输出严格满足概率和为 1。函数签名与关键参数见 root_finding.pyentmax_bisect(X, alpha1.5, dim-1, n_iter50, ensure_sum_oneTrue, trainingTrue) sparsemax_bisect(X, dim-1, n_iter50, ensure_sum_oneTrue, trainingTrue)alphaα 取值1。alpha2等价于 sparsemax理论上alpha1会退化为 softmax但数值上不支持 α1需要 softmax 时请直接使用torch.nn.functional.softmax源码 docstring 对此有明确警告n_iter二分迭代次数float32 精度下约24 次即可达到机器精度默认 50 次留足余量ensure_sum_one是否对结果做归一化保证输出精确满足Σp 1防止下游数值问题alpha还支持传入torch.Tensor可以为每个样本行指定不同的 α形状需满足沿dim维度为 1、其余维度与 X 一致或可广播。此外还提供了 sparsemax 的二分特例SparsemaxBisectFunction通过重写_gpx ** (α-1)与_gp_invy ** (1/(α-1))为恒等映射省去幂运算开销见 root_finding.py 第 106-141 行。4.2 对 α 求梯度自适应稀疏度entmax_bisect是同时关于 X 和 α 可微的。EntmaxBisectFunction.backward在ctx.needs_input_grad[1]为真时利用香农熵项S Y * log(Y)与偏斜分布Y_skewed推导出d_alpha的闭式表达式见 root_finding.py 第 48-65 行注释特别提醒α 不能接近 1否则存在不确定性indetermination。这意味着我们可以把 α 当作可学习参数让模型自己决定注意力应该多稀疏。README 给出的对 α 求梯度的完整示例import torch from torch.autograd import grad from ludwig.utils.entmax import entmax_bisect x torch.tensor([[-1, 0, 0.5], [1, 2, 3.5]]) alpha torch.tensor(1.33, requires_gradTrue) p entmax_bisect(x, alpha) print(p) # tensor([[0.0460, 0.3276, 0.6264], # [0.0026, 0.1012, 0.8963]], grad_fnEntmaxBisectFunctionBackward) print(grad(p[0, 0], alpha)) # (tensor(-0.2562),)输出张量的grad_fn为EntmaxBisectFunctionBackward且p[0,0]对alpha的梯度可正常回传——这正是 Gradients w.r.t. alpha for adaptive, learned sparsity 这一特性的落地方式也是后续 TabNet 中adaptive模式的数学基础。五、深入实现partial-sort 算法与 k 参数对于 1.5-entmax 与 sparsemax 两个特例activations.py 采用精确的 partial-sort部分排序算法避免对整行做全量排序从而获得更优的常数因子性能。_sparsemax_threshold_and_support与_entmax_threshold_and_support都接受一个k参数kNone或k ≥ X.shape[dim]直接对全行排序torch.sort(..., descendingTrue)k X.shape[dim]改用torch.topk(X, kk, dimdim)只取最大的 k 个元素期望稀疏解的非零个数不超过 k递归兜底如果解的实际稀疏度超过 ksupport_size k说明支持集被截断算法会以2*k的倍增计划递归调用自身直到求出完整的支持集见_roll_last 递归分支。在 README 与 activations.py 的 docstring 中都强调k的取值应略大于预期的非零个数以获得最优性能。k同时作为函数式接口与Sparsemax/Entmax15模块的构造参数对外暴露。5.1 自定义反向传播的推导要点SparsemaxFunction.backward与Entmax15Function.backward都只依赖正向保存的输出张量即可完成梯度计算无需保存中间排序结果sparsemax对输出为 0 的位置梯度置零再对支持集内做去均值投影grad_input - v_hat这正是单纯形投影的雅可比形式entmax15利用gppr Y.sqrt()即1 / g(Y)计算加权修正实现dX - q * gppr的投影。两者都返回dX, None, None即只对输入 X 求梯度。六、训练与推理分离torchscript 兼容设计sparsemax、entmax15、entmax_bisect、sparsemax_bisect四个函数都带有trainingTrue关键字参数源码注释明确说明了原因自定义torch.autograd.Function无法被 torchscript 序列化activations.py 第 229-234 行引用了 PyTorch issue #22329。因此训练阶段trainingTrue走CustomFunction.apply(...)路径获得自定义反向传播与对 α 的梯度推理/导出阶段trainingFalse直接调用纯 PyTorch 算子组合的正向函数_sparsemax_forward/_entmax15_forward/_entmax_bisect_forward保证模型可被 torchscript 追踪导出。对应的nn.Module封装Sparsemax、Entmax15、EntmaxBisect、SparsemaxBisect在forward中传入self.training因此用户在模型eval()/train()切换时会自动选择正确路径无需手动干预。七、稀疏损失函数Fenchel-Young 损失家族entmax 不仅提供了稀疏概率映射还配套了对应的稀疏损失函数——它们是 softmax 交叉熵的稀疏替代品通过 Fenchel 对偶推导而来天然与对应激活函数匹配。全部实现位于 losses.pyREADME 与 loss_modules.py 的注释均将其称为激活函数 Ω 函数的 Fenchel 共轭损失。函数式接口一览函数对应映射特点sparsemax_loss(X, target, kNone)sparsemaxpartial-sort 计算ω (1 - Σp²) / 2sparsemax_bisect_loss(X, target, n_iter50)sparsemax二分法计算entmax15_loss(X, target, kNone)1.5-entmaxpartial-sort 计算ω (1 - Σ p·√p) / 0.75entmax_bisect_loss(X, target, alpha1.5, n_iter50)α-entmax二分法计算ω (1 - Σp^α) / (α(α-1))支持逐行 α它们的输入约定与交叉熵一致X为形状(n_samples, n_classes)的分数logitstarget为形状(n_samples,)的长整型标签满足0 ≤ target n_classes。_GenericLossFunction的正向传播非常优雅先用对应激活函数求出最优概率p_star计算正则项ω(p_star, α)再通过scatter_add_把目标类位置的贡献扣除并加上线性项einsum(ij,ij-i, p_star, X)最终得到逐样本损失反向传播只需保存p_star即可见 losses.py 第 39-64 行。除了函数式接口还提供了带ignore_index与reduction的nn.Module封装类SparsemaxLoss、SparsemaxBisectLoss、Entmax15Loss、EntmaxBisectLoss。它们继承自_GenericLossignore_index默认值为 Ludwig 的IGNORE_INDEX_TOKEN_ID可屏蔽填充位置如序列任务中的 pad token被屏蔽位置的损失被masked_fill_置零reduction支持elementwise_mean默认等价于求均值、sum、none三种模式非法值会抛出带提示的ValueError当所有行都被 ignore 时返回零损失与零梯度避免除零见 losses.py 第 31-34 行。from ludwig.utils.entmax.losses import Entmax15Loss criterion Entmax15Loss(k100) # 默认 k100可配合 ignore_index 使用 loss criterion(logits, targets) # logits: [B, C], targets: [B]八、在 Ludwig 中的实际集成损失模块与 TabNet8.1 注册为输出特征的损失函数Ludwig 的损失模块 loss_modules.py 将SparsemaxLoss与Entmax15Loss注册为可用损失类型通过register_loss装饰器对应SparsemaxLossConfig与Entmax15LossConfig。其 docstring 明确阐述了定位SparsemaxLosssoftmax 交叉熵的稀疏替代品……对 sparsemax 支持集之外的类别分配零梯度在概率单纯形上产生精确稀疏性参考文献为 Martins Astudillo (ICML 2016)Entmax15Lossα1.5 entmax 激活的 Fenchel 共轭损失……产生介于 softmax稠密与 sparsemax最大稀疏之间的中等稀疏概率分布参考文献为 Peters et al. (ACL 2019)。这两个损失类都实现了LogitsInputsMixin直接接收[B, C]的 logits 与[B]的长整型标签内部委托给ludwig.utils.entmax.losses中的底层实现。这意味着你可以在 Ludwig 配置中为分类输出特征直接指定type: sparsemax或type: entmax15作为损失函数无需编写任何自定义代码。8.2 TabNet 组合器中的 entmax 稀疏注意力entmax 在 Ludwig 中最引人注目的应用是TabNet 组合器的注意力模块。tabnet_modules.py 从ludwig.utils.entmax导入了Entmax15、EntmaxBisect、Sparsemax在AttentiveTransformer中依据entmax_mode选择稀疏映射if self.entmax_mode sparsemax: self.entmax_module Sparsemax() elif self.entmax_mode entmax15: self.entmax_module Entmax15() else: # adaptive self.entmax_module EntmaxBisect(alphaself.trainable_alpha)其中adaptive模式把 α 注册为可学习缓冲区self.register_buffer(trainable_alpha, torch.tensor(entmax_alpha, requires_gradTrue))正是利用了本文第四节介绍的对 α 求梯度能力让模型在训练中自动调节特征选择的稀疏程度。对应的配置 schema 位于 tabnet.py两个参数的约束如下entmax_modeStringOptions可选值[entmax15, sparsemax, constant, adaptive]默认sparsemax。文档说明entmax 是一族推广 softmax 与 sparsemax 的稀疏概率映射entmax_mode控制稀疏度entmax_alphaFloatRange取值范围1.0 到 2.0默认 1.5。schema 注释明确1对应 softmax2对应 sparsemax当entmax_mode为adaptive时它作为可学习参数的初始值。在 TabNet 配置文件中即可直接使用combiner: type: tabnet entmax_mode: adaptive entmax_alpha: 1.5九、测试验证数值一致性与梯度正确性仓库为 entmax 提供了完善的单元测试位于 tests/ludwig/utils/entmax/test_root_finding.py验证partial-sort 实现与二分法实现输出一致如test_sparsemax断言sparsemax(x, 1)与sparsemax_bisect(x)的平方误差小于1e-7test_entmax15同理并在训练/推理两种模式下都做了验证同时使用torch.autograd.gradcheck对sparsemax_bisect做梯度数值校验test_sparsemax_grad确认自定义反向传播的正确性test_losses.py验证各稀疏损失函数的数值行为test_topk.py 与 test_mask.py分别验证 partial-sort 的k参数路径与ignore_index掩码逻辑。此外test_tabnet_modules.py 覆盖了 TabNet 组合器中 entmax 相关参数的组合。这些测试既保证了数学实现的正确性也保证了集成进 Ludwig 后各配置组合的行为稳定。十、论文出处与进一步阅读entmax 的理论根基来自两篇核心论文原文以 BibTeX 形式记录于 READMEPeters, Niculae, and Martins, Sparse Sequence-to-Sequence Models, Proc. ACL 2019—— 提出 1.5-entmax 与通用 α-entmax并将其应用于序列到序列模型的注意力稀疏化Correia, Niculae, and Martins, Adaptively Sparse Transformers, Proc. EMNLP-IJCNLP 2019—— 提出对 α 的梯度计算实现自适应稀疏 Transformer。如需进一步理解README 还推荐了三篇延伸文献Blondel, Martins, and Niculae, 2019, Learning with Fenchel-Young Losses——本仓库损失函数的理论框架来源Martins and Astudillo, 2016, From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification——sparsemax 的原始出处Peters and Martins, 2019, IT-IST at the SIGMORPHON 2019 Shared Task: Sparse Two-headed Models for Inflection——entmax 在词形变化任务上的工程应用。结语entmax 家族为深度学习提供了一条从稠密到稀疏的连续调节路径sparsemax 给出最激进的精确稀疏entmax15 提供折中的半稀疏而基于二分法的 α-entmax 配合对 α 的可微性让稀疏度本身成为可学习的模型参数。在 Ludwig 中这套实现不仅以独立子包形式完整保留ludwig/utils/entmax/还深度集成了输出特征的稀疏损失sparsemax/entmax15与 TabNet 组合器的稀疏特征选择entmax_mode: sparsemax | entmax15 | adaptive。无论你是想在自定义模型中引入稀疏注意力还是在 Ludwig 配置中直接启用稀疏损失本文涉及的 API、参数与源码路径都可作为直接的实践依据。【免费下载链接】ludwigLow-code framework for building custom LLMs, neural networks, and other AI models项目地址: https://gitcode.com/gh_mirrors/lu/ludwig创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考