Ludwig 中的 entmax從 sparsemax 到可學(xué)習(xí)稀疏性的 α-entmax 概率映射與損失函數(shù)實(shí)戰(zhàn)指南【免費(fèi)下載鏈接】ludwigLow-code framework for building custom LLMs, neural networks, and other AI models項(xiàng)目地址: https://gitcode.com/gh_mirrors/lu/ludwig導(dǎo)讀本文圍繞 Ludwig 倉庫內(nèi)置的 entmax 實(shí)現(xiàn) 展開系統(tǒng)講解一類稀疏概率映射及其配套損失函數(shù)entmax 家族以 Tsallis α-熵推廣 softmax在保留可微性的同時(shí)讓輸出概率產(chǎn)生精確的零值從而在注意力、多標(biāo)簽分類、TabNet 特征選擇等場(chǎng)景中天然具備稀疏性與可解釋性。讀完本文你將掌握sparsemax、entmax15、通用α-entmax含對(duì) α 的可學(xué)習(xí)梯度的數(shù)學(xué)原理、PyTorch API 用法、損失函數(shù)設(shè)計(jì)以及它們?nèi)绾卧?Ludwig 的損失模塊與 TabNet 組合器中落地配置。一、entmax 是什么稀疏版的 softmax 概率映射傳統(tǒng) softmax 將任意實(shí)數(shù)向量映射為嚴(yán)格正的、處處非零的概率分布而 entmax 是一族稀疏概率映射sparse probability mapping它把輸入向量映射到單純形simplex上并允許部分分量精確等于 0。這種稀疏性對(duì)注意力機(jī)制、多標(biāo)簽分類、序列到序列模型尤其有價(jià)值模型可以顯式地忽略無關(guān)信息而不是分配給它們一個(gè)微小但非零的權(quán)重。正如 entmax README 所概括該包提供精確的 partial-sort 算法針對(duì) 1.5-entmax 與 2-entmax即 sparsemax兩種特例基于二分法bisection的通用算法支持任意 α-entmax對(duì) α 的梯度實(shí)現(xiàn)自適應(yīng)、可學(xué)習(xí)的稀疏度控制。在數(shù)學(xué)上α-entmax 求解如下帶約束的優(yōu)化問題見 root_finding.py 源碼注釋max_p x, p - H_α(p) s.t. p ≥ 0, Σp 1其中H_α(p)是 Tsallis α-熵。α 的不同取值構(gòu)成一條連續(xù)譜系α 取值對(duì)應(yīng)映射稀疏程度α → 1softmax無稀疏性處處非零α 1.51.5-entmaxentmax15中等稀疏介于兩者之間α 22-entmax即sparsemax最大程度稀疏等價(jià)于歐氏投影到單純形sparsemax 的幾何意義非常直觀它是把向量 x 投影到概率單純形上的歐氏投影min_p ||x - p||?投影結(jié)果天然具有精確零分量。這一思想源自 Martins Astudillo2016隨后由 Peters 等人2019推廣為 1.5-entmax 與通用 α-entmax并用于稀疏序列到序列模型。二、安裝與導(dǎo)入倉庫內(nèi)嵌的 entmax 包entmax 在 Ludwig 中以內(nèi)部子包的形式存在于ludwig/utils/entmax/目錄共五個(gè)文件init.py版本號(hào)與全部公開 API 的導(dǎo)出activations.pysparsemax、entmax15 的函數(shù)式實(shí)現(xiàn)與nn.Module封裝root_finding.py通用 α-entmax 的二分法實(shí)現(xiàn)含對(duì) α 的反向傳播losses.py對(duì)應(yīng)的稀疏損失函數(shù)與損失類README.md 與 LICENSEMIT 協(xié)議版權(quán)歸屬 DeepSPIN2019。獨(dú)立使用時(shí)可以通過 PyPI 安裝原版包pip install entmax而在 Ludwig 項(xiàng)目中只需從內(nèi)置路徑導(dǎo)入無需額外安裝例如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__涵蓋函數(shù)式接口小寫命名、損失函數(shù)與對(duì)應(yīng)的nn.Module類大寫命名。運(yùn)行環(huán)境要求為 Python 3、PyTorch ≥ 1.0單元測(cè)試需要 pytest。三、快速上手sparsemax 與 entmax15 的直觀對(duì)比entmax README 給出了一個(gè)非常直觀的三路對(duì)比示例這里完整復(fù)現(xiàn)并補(bǔ)充解讀。對(duì)輸入向量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給三個(gè)位置都分配了非零概率0.0486 / 0.3592 / 0.5922即使-2明顯是無關(guān)項(xiàng)sparsemaxα2把-2對(duì)應(yīng)位置的概率精確置為 0其余兩項(xiàng)按 0.25 / 0.75 分配entmax15α1.5同樣把-2置零但概率分配更平滑0.326 / 0.674稀疏性介于 softmax 與 sparsemax 之間。從 activations.py 源碼 可以看到兩者的計(jì)算內(nèi)核_sparsemax_threshold_and_support通過累計(jì)和與閾值公式計(jì)算最優(yōu)閾值τ和支持集大小非零分量個(gè)數(shù)輸出clamp(X - τ, min0)_entmax_threshold_and_support基于均值、均方與方差公式ss rho * (mean_sq - mean**2)delta (1 - ss) / rho計(jì)算 1.5-entmax 的閾值輸出clamp(X - τ, min0) ** 2因此 entmax15 的輸出是閾值的平方。兩者的正向傳播都采用了與 softmax 相同的數(shù)值穩(wěn)定性技巧——先減去每行最大值X X - max_val避免大數(shù)值溢出。3.1 函數(shù)式接口與 nn.Module 封裝sparsemax與entmax15同時(shí)提供函數(shù)式接口與模塊化封裝見 activations.pyfrom ludwig.utils.entmax import Sparsemax, Entmax15 # 函數(shù)式直接作用于張量 p sparsemax(x, dim-1) # 模塊化便于嵌入 nn.Sequential 或自定義網(wǎng)絡(luò)層 layer Sparsemax(dim-1, kNone) # kNone 表示全排序 p layer(x) layer Entmax15(dim-1, kNone) p layer(x)兩者的nn.Module版本Sparsemax、Entmax15都接受dim沿哪個(gè)維度歸一化與kpartial-sort 的元素?cái)?shù)量詳見第五節(jié)兩個(gè)構(gòu)造參數(shù)并在forward中透?jìng)鱯elf.training以保證推理階段走 torchscript 兼容的快速路徑詳見第六節(jié)。四、通用 α-entmax二分法實(shí)現(xiàn)與可學(xué)習(xí)稀疏度4.1 支持任意 α 的二分求根算法當(dāng) α 不是 1.5 或 2 這樣的特例時(shí)無法使用解析的 partial-sort 公式root_finding.py提供了基于**二分法bisection即二分搜索求根**的通用算法entmax_bisect。其核心思想是對(duì)每個(gè)向量在[τ_lo, τ_hi]區(qū)間內(nèi)反復(fù)二分尋找使Σ p(x_i - τ) 1成立的閾值τ迭代n_iter次后得到近似投影最后通過ensure_sum_one歸一化保證輸出嚴(yán)格滿足概率和為 1。函數(shù)簽名與關(guān)鍵參數(shù)見 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等價(jià)于 sparsemax理論上alpha1會(huì)退化為 softmax但數(shù)值上不支持 α1需要 softmax 時(shí)請(qǐng)直接使用torch.nn.functional.softmax源碼 docstring 對(duì)此有明確警告n_iter二分迭代次數(shù)float32 精度下約24 次即可達(dá)到機(jī)器精度默認(rèn) 50 次留足余量ensure_sum_one是否對(duì)結(jié)果做歸一化保證輸出精確滿足Σp 1防止下游數(shù)值問題alpha還支持傳入torch.Tensor可以為每個(gè)樣本行指定不同的 α形狀需滿足沿dim維度為 1、其余維度與 X 一致或可廣播。此外還提供了 sparsemax 的二分特例SparsemaxBisectFunction通過重寫_gpx ** (α-1)與_gp_invy ** (1/(α-1))為恒等映射省去冪運(yùn)算開銷見 root_finding.py 第 106-141 行。4.2 對(duì) α 求梯度自適應(yīng)稀疏度entmax_bisect是同時(shí)關(guān)于 X 和 α 可微的。EntmaxBisectFunction.backward在ctx.needs_input_grad[1]為真時(shí)利用香農(nóng)熵項(xiàng)S Y * log(Y)與偏斜分布Y_skewed推導(dǎo)出d_alpha的閉式表達(dá)式見 root_finding.py 第 48-65 行注釋特別提醒α 不能接近 1否則存在不確定性indetermination。這意味著我們可以把 α 當(dāng)作可學(xué)習(xí)參數(shù)讓模型自己決定注意力應(yīng)該多稀疏。README 給出的對(duì) α 求梯度的完整示例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]對(duì)alpha的梯度可正?；貍鳌@正是 Gradients w.r.t. alpha for adaptive, learned sparsity 這一特性的落地方式也是后續(xù) TabNet 中adaptive模式的數(shù)學(xué)基礎(chǔ)。五、深入實(shí)現(xiàn)partial-sort 算法與 k 參數(shù)對(duì)于 1.5-entmax 與 sparsemax 兩個(gè)特例activations.py 采用精確的 partial-sort部分排序算法避免對(duì)整行做全量排序從而獲得更優(yōu)的常數(shù)因子性能。_sparsemax_threshold_and_support與_entmax_threshold_and_support都接受一個(gè)k參數(shù)kNone或k ≥ X.shape[dim]直接對(duì)全行排序torch.sort(..., descendingTrue)k X.shape[dim]改用torch.topk(X, kk, dimdim)只取最大的 k 個(gè)元素期望稀疏解的非零個(gè)數(shù)不超過 k遞歸兜底如果解的實(shí)際稀疏度超過 ksupport_size k說明支持集被截?cái)嗨惴〞?huì)以2*k的倍增計(jì)劃遞歸調(diào)用自身直到求出完整的支持集見_roll_last 遞歸分支。在 README 與 activations.py 的 docstring 中都強(qiáng)調(diào)k的取值應(yīng)略大于預(yù)期的非零個(gè)數(shù)以獲得最優(yōu)性能。k同時(shí)作為函數(shù)式接口與Sparsemax/Entmax15模塊的構(gòu)造參數(shù)對(duì)外暴露。5.1 自定義反向傳播的推導(dǎo)要點(diǎn)SparsemaxFunction.backward與Entmax15Function.backward都只依賴正向保存的輸出張量即可完成梯度計(jì)算無需保存中間排序結(jié)果sparsemax對(duì)輸出為 0 的位置梯度置零再對(duì)支持集內(nèi)做去均值投影grad_input - v_hat這正是單純形投影的雅可比形式entmax15利用gppr Y.sqrt()即1 / g(Y)計(jì)算加權(quán)修正實(shí)現(xiàn)dX - q * gppr的投影。兩者都返回dX, None, None即只對(duì)輸入 X 求梯度。六、訓(xùn)練與推理分離torchscript 兼容設(shè)計(jì)sparsemax、entmax15、entmax_bisect、sparsemax_bisect四個(gè)函數(shù)都帶有trainingTrue關(guān)鍵字參數(shù)源碼注釋明確說明了原因自定義torch.autograd.Function無法被 torchscript 序列化activations.py 第 229-234 行引用了 PyTorch issue #22329。因此訓(xùn)練階段trainingTrue走CustomFunction.apply(...)路徑獲得自定義反向傳播與對(duì) α 的梯度推理/導(dǎo)出階段trainingFalse直接調(diào)用純 PyTorch 算子組合的正向函數(shù)_sparsemax_forward/_entmax15_forward/_entmax_bisect_forward保證模型可被 torchscript 追蹤導(dǎo)出。對(duì)應(yīng)的nn.Module封裝Sparsemax、Entmax15、EntmaxBisect、SparsemaxBisect在forward中傳入self.training因此用戶在模型eval()/train()切換時(shí)會(huì)自動(dòng)選擇正確路徑無需手動(dòng)干預(yù)。七、稀疏損失函數(shù)Fenchel-Young 損失家族e(cuò)ntmax 不僅提供了稀疏概率映射還配套了對(duì)應(yīng)的稀疏損失函數(shù)——它們是 softmax 交叉熵的稀疏替代品通過 Fenchel 對(duì)偶推導(dǎo)而來天然與對(duì)應(yīng)激活函數(shù)匹配。全部實(shí)現(xiàn)位于 losses.pyREADME 與 loss_modules.py 的注釋均將其稱為激活函數(shù) Ω 函數(shù)的 Fenchel 共軛損失。函數(shù)式接口一覽函數(shù)對(duì)應(yīng)映射特點(diǎn)sparsemax_loss(X, target, kNone)sparsemaxpartial-sort 計(jì)算ω (1 - Σp2) / 2sparsemax_bisect_loss(X, target, n_iter50)sparsemax二分法計(jì)算entmax15_loss(X, target, kNone)1.5-entmaxpartial-sort 計(jì)算ω (1 - Σ p·√p) / 0.75entmax_bisect_loss(X, target, alpha1.5, n_iter50)α-entmax二分法計(jì)算ω (1 - Σp^α) / (α(α-1))支持逐行 α它們的輸入約定與交叉熵一致X為形狀(n_samples, n_classes)的分?jǐn)?shù)logitstarget為形狀(n_samples,)的長(zhǎng)整型標(biāo)簽滿足0 ≤ target n_classes。_GenericLossFunction的正向傳播非常優(yōu)雅先用對(duì)應(yīng)激活函數(shù)求出最優(yōu)概率p_star計(jì)算正則項(xiàng)ω(p_star, α)再通過scatter_add_把目標(biāo)類位置的貢獻(xiàn)扣除并加上線性項(xiàng)einsum(ij,ij-i, p_star, X)最終得到逐樣本損失反向傳播只需保存p_star即可見 losses.py 第 39-64 行。除了函數(shù)式接口還提供了帶ignore_index與reduction的nn.Module封裝類SparsemaxLoss、SparsemaxBisectLoss、Entmax15Loss、EntmaxBisectLoss。它們繼承自_GenericLossignore_index默認(rèn)值為 Ludwig 的IGNORE_INDEX_TOKEN_ID可屏蔽填充位置如序列任務(wù)中的 pad token被屏蔽位置的損失被masked_fill_置零reduction支持elementwise_mean默認(rèn)等價(jià)于求均值、sum、none三種模式非法值會(huì)拋出帶提示的ValueError當(dāng)所有行都被 ignore 時(shí)返回零損失與零梯度避免除零見 losses.py 第 31-34 行。from ludwig.utils.entmax.losses import Entmax15Loss criterion Entmax15Loss(k100) # 默認(rèn) k100可配合 ignore_index 使用 loss criterion(logits, targets) # logits: [B, C], targets: [B]八、在 Ludwig 中的實(shí)際集成損失模塊與 TabNet8.1 注冊(cè)為輸出特征的損失函數(shù)Ludwig 的損失模塊 loss_modules.py 將SparsemaxLoss與Entmax15Loss注冊(cè)為可用損失類型通過register_loss裝飾器對(duì)應(yīng)SparsemaxLossConfig與Entmax15LossConfig。其 docstring 明確闡述了定位SparsemaxLosssoftmax 交叉熵的稀疏替代品……對(duì) sparsemax 支持集之外的類別分配零梯度在概率單純形上產(chǎn)生精確稀疏性參考文獻(xiàn)為 Martins Astudillo (ICML 2016)Entmax15Lossα1.5 entmax 激活的 Fenchel 共軛損失……產(chǎn)生介于 softmax稠密與 sparsemax最大稀疏之間的中等稀疏概率分布參考文獻(xiàn)為 Peters et al. (ACL 2019)。這兩個(gè)損失類都實(shí)現(xiàn)了LogitsInputsMixin直接接收[B, C]的 logits 與[B]的長(zhǎng)整型標(biāo)簽內(nèi)部委托給ludwig.utils.entmax.losses中的底層實(shí)現(xiàn)。這意味著你可以在 Ludwig 配置中為分類輸出特征直接指定type: sparsemax或type: entmax15作為損失函數(shù)無需編寫任何自定義代碼。8.2 TabNet 組合器中的 entmax 稀疏注意力entmax 在 Ludwig 中最引人注目的應(yīng)用是TabNet 組合器的注意力模塊。tabnet_modules.py 從ludwig.utils.entmax導(dǎo)入了Entmax15、EntmaxBisect、Sparsemax在AttentiveTransformer中依據(jù)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模式把 α 注冊(cè)為可學(xué)習(xí)緩沖區(qū)self.register_buffer(trainable_alpha, torch.tensor(entmax_alpha, requires_gradTrue))正是利用了本文第四節(jié)介紹的對(duì) α 求梯度能力讓模型在訓(xùn)練中自動(dòng)調(diào)節(jié)特征選擇的稀疏程度。對(duì)應(yīng)的配置 schema 位于 tabnet.py兩個(gè)參數(shù)的約束如下entmax_modeStringOptions可選值[entmax15, sparsemax, constant, adaptive]默認(rèn)sparsemax。文檔說明entmax 是一族推廣 softmax 與 sparsemax 的稀疏概率映射entmax_mode控制稀疏度entmax_alphaFloatRange取值范圍1.0 到 2.0默認(rèn) 1.5。schema 注釋明確1對(duì)應(yīng) softmax2對(duì)應(yīng) sparsemax當(dāng)entmax_mode為adaptive時(shí)它作為可學(xué)習(xí)參數(shù)的初始值。在 TabNet 配置文件中即可直接使用combiner: type: tabnet entmax_mode: adaptive entmax_alpha: 1.5九、測(cè)試驗(yàn)證數(shù)值一致性與梯度正確性倉庫為 entmax 提供了完善的單元測(cè)試位于 tests/ludwig/utils/entmax/test_root_finding.py驗(yàn)證partial-sort 實(shí)現(xiàn)與二分法實(shí)現(xiàn)輸出一致如test_sparsemax斷言sparsemax(x, 1)與sparsemax_bisect(x)的平方誤差小于1e-7test_entmax15同理并在訓(xùn)練/推理兩種模式下都做了驗(yàn)證同時(shí)使用torch.autograd.gradcheck對(duì)sparsemax_bisect做梯度數(shù)值校驗(yàn)test_sparsemax_grad確認(rèn)自定義反向傳播的正確性test_losses.py驗(yàn)證各稀疏損失函數(shù)的數(shù)值行為test_topk.py 與 test_mask.py分別驗(yàn)證 partial-sort 的k參數(shù)路徑與ignore_index掩碼邏輯。此外test_tabnet_modules.py 覆蓋了 TabNet 組合器中 entmax 相關(guān)參數(shù)的組合。這些測(cè)試既保證了數(shù)學(xué)實(shí)現(xiàn)的正確性也保證了集成進(jìn) Ludwig 后各配置組合的行為穩(wěn)定。十、論文出處與進(jìn)一步閱讀entmax 的理論根基來自兩篇核心論文原文以 BibTeX 形式記錄于 READMEPeters, Niculae, and Martins, Sparse Sequence-to-Sequence Models, Proc. ACL 2019—— 提出 1.5-entmax 與通用 α-entmax并將其應(yīng)用于序列到序列模型的注意力稀疏化Correia, Niculae, and Martins, Adaptively Sparse Transformers, Proc. EMNLP-IJCNLP 2019—— 提出對(duì) α 的梯度計(jì)算實(shí)現(xiàn)自適應(yīng)稀疏 Transformer。如需進(jìn)一步理解README 還推薦了三篇延伸文獻(xiàn)Blondel, Martins, and Niculae, 2019, Learning with Fenchel-Young Losses——本倉庫損失函數(shù)的理論框架來源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 在詞形變化任務(wù)上的工程應(yīng)用。結(jié)語entmax 家族為深度學(xué)習(xí)提供了一條從稠密到稀疏的連續(xù)調(diào)節(jié)路徑sparsemax 給出最激進(jìn)的精確稀疏entmax15 提供折中的半稀疏而基于二分法的 α-entmax 配合對(duì) α 的可微性讓稀疏度本身成為可學(xué)習(xí)的模型參數(shù)。在 Ludwig 中這套實(shí)現(xiàn)不僅以獨(dú)立子包形式完整保留ludwig/utils/entmax/還深度集成了輸出特征的稀疏損失sparsemax/entmax15與 TabNet 組合器的稀疏特征選擇entmax_mode: sparsemax | entmax15 | adaptive。無論你是想在自定義模型中引入稀疏注意力還是在 Ludwig 配置中直接啟用稀疏損失本文涉及的 API、參數(shù)與源碼路徑都可作為直接的實(shí)踐依據(jù)?！久赓M(fèi)下載鏈接】ludwigLow-code framework for building custom LLMs, neural networks, and other AI models項(xiàng)目地址: https://gitcode.com/gh_mirrors/lu/ludwig創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考