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