企业官网建设流程全解析

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 α-熵。α 的不同取值构成一条连续谱系：

sparsemax 的几何意义非常直观：它是把向量 x 投影到概率单纯形上的欧氏投影（min_p ||x - p||₂），投影结果天然具有精确零分量。这一思想源自 Martins & Astudillo（2016），随后由 Peters 等人（2019）推广为 1.5-entmax 与通用 α-entmax，并用于稀疏序列到序列模型。

二、安装与导入：仓库内嵌的 entmax 包

entmax 在 Ludwig 中以内部子包的形式存在于ludwig/utils/entmax/目录，共五个文件：

• init.py：版本号与全部公开 API 的导出；
• activations.py：sparsemax、entmax15 的函数式实现与nn.Module封装；
• root_finding.py：通用 α-entmax 的二分法实现，含对 α 的反向传播；
• losses.py：对应的稀疏损失函数与损失类；
• README.md 与 LICENSE（MIT 协议，版权归属 DeepSPIN，2019）。

独立使用时，可以通过 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 与 entmax15：

import torch from torch.nn.functional import softmax from ludwig.utils.entmax import sparsemax, entmax15 x = torch.tensor([-2, 0, 0.5]) print(softmax(x, dim=0)) # tensor([0.0486, 0.3592, 0.5922]) print(sparsemax(x, dim=0)) # tensor([0.0000, 0.2500, 0.7500]) print(entmax15(x, dim=0)) # 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 - τ, min=0)；
• _entmax_threshold_and_support基于均值、均方与方差公式（ss = rho * (mean_sq - mean**2)，delta = (1 - ss) / rho）计算 1.5-entmax 的阈值，输出clamp(X - τ, min=0) ** 2，因此 entmax15 的输出是"阈值的平方"。

两者的正向传播都采用了与 softmax 相同的数值稳定性技巧——先减去每行最大值（X = X - max_val），避免大数值溢出。

3.1 函数式接口与 nn.Module 封装

sparsemax与entmax15同时提供函数式接口与模块化封装（见 activations.py）：

from ludwig.utils.entmax import Sparsemax, Entmax15 # 函数式：直接作用于张量 p = sparsemax(x, dim=-1) # 模块化：便于嵌入 nn.Sequential 或自定义网络层 layer = Sparsemax(dim=-1, k=None) # k=None 表示全排序 p = layer(x) layer = Entmax15(dim=-1, k=None) p = layer(x)

两者的nn.Module版本（Sparsemax、Entmax15）都接受dim（沿哪个维度归一化）与k（partial-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.py）：

entmax_bisect(X, alpha=1.5, dim=-1, n_iter=50, ensure_sum_one=True, training=True) sparsemax_bisect(X, dim=-1, n_iter=50, ensure_sum_one=True, training=True)

• alpha：α 取值（>1）。alpha=2等价于 sparsemax；理论上alpha=1会退化为 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，通过重写_gp（x ** (α-1)）与_gp_inv（y ** (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_grad=True) p = entmax_bisect(x, alpha) print(p) # tensor([[0.0460, 0.3276, 0.6264], # [0.0026, 0.1012, 0.8963]], grad_fn=<EntmaxBisectFunctionBackward>) 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参数：

• k=None或k ≥ X.shape[dim]：直接对全行排序（torch.sort(..., descending=True)）；
• k < X.shape[dim]：改用torch.topk(X, k=k, dim=dim)只取最大的 k 个元素，期望稀疏解的非零个数不超过 k；
• 递归兜底：如果解的实际稀疏度超过 k（support_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四个函数都带有training=True关键字参数，源码注释明确说明了原因：自定义torch.autograd.Function无法被 torchscript 序列化（activations.py 第 229-234 行引用了 PyTorch issue #22329）。

因此：

• 训练阶段（training=True）：走CustomFunction.apply(...)路径，获得自定义反向传播与对 α 的梯度；
• 推理/导出阶段（training=False）：直接调用纯 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.py，README 与 loss_modules.py 的注释均将其称为"激活函数 Ω 函数的 Fenchel 共轭损失"。

函数式接口一览：

它们的输入约定与交叉熵一致：X为形状(n_samples, n_classes)的分数（logits），target为形状(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。它们继承自_GenericLoss：

• ignore_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(k=100) # 默认 k=100，可配合 ignore_index 使用 loss = criterion(logits, targets) # logits: [B, C], targets: [B]

八、在 Ludwig 中的实际集成：损失模块与 TabNet

8.1 注册为输出特征的损失函数

Ludwig 的损失模块 loss_modules.py 将SparsemaxLoss与Entmax15Loss注册为可用损失类型（通过@register_loss装饰器，对应SparsemaxLossConfig与Entmax15LossConfig）。其 docstring 明确阐述了定位：

• SparsemaxLoss："softmax 交叉熵的稀疏替代品……对 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(alpha=self.trainable_alpha)

其中adaptive模式把 α 注册为可学习缓冲区（self.register_buffer("trainable_alpha", torch.tensor(entmax_alpha, requires_grad=True))），正是利用了本文第四节介绍的"对 α 求梯度"能力，让模型在训练中自动调节特征选择的稀疏程度。

对应的配置 schema 位于 tabnet.py，两个参数的约束如下：

• entmax_mode：StringOptions，可选值["entmax15", "sparsemax", "constant", "adaptive"]，默认"sparsemax"。文档说明"entmax 是一族推广 softmax 与 sparsemax 的稀疏概率映射，entmax_mode控制稀疏度"；
• entmax_alpha：FloatRange，取值范围1.0 到 2.0，默认 1.5。schema 注释明确：1对应 softmax，2对应 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-7；test_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 形式记录于 README）：

1. Peters, Niculae, and Martins, "Sparse Sequence-to-Sequence Models", Proc. ACL 2019—— 提出 1.5-entmax 与通用 α-entmax，并将其应用于序列到序列模型的注意力稀疏化；
2. 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