技术摘要：超越 Softmax：分类随机变量的自然参数化

问题表述

潜在分类随机变量是深度学习架构中的基本组成部分，用于建模强化学习（RL）中的离散动作、在潜在变量模型中表示主题，以及在图神经网络（GNN）中定义图结构。然而，这些变量的离散性质给梯度下降学习算法带来了重大挑战。虽然训练信号通常源自下游任务（例如变分自编码器中的重构损失或强化学习中的奖励），但缺乏针对潜在变量的直接监督信号，因此必须使用梯度估计器。

标准方法通常依赖得分函数（REINFORCE）估计器，该估计器无偏但方差较高，导致训练不稳定。替代方法（如 Gumbel-Softmax 技巧）引入连续松弛以支持反向传播，但可能会引入偏差。本文探讨的互补视角聚焦于分类分布本身的参数化。作者认为，无处不在的softmax函数虽然标准，但会诱导一个稠密的费雪信息矩阵（FIM）。这种稠密性在参数空间中造成几何畸变，导致基于梯度的优化器遵循振荡且不够直接的轨迹，从而阻碍收敛和稳定性。

方法论：Catnat 参数化

本文提出了catnat，这是一种旨在与信息几何原理保持一致的分类随机变量新参数化方法。

1. 信息几何视角

作者利用费雪信息矩阵（FIM）分析了参数空间的几何结构。对于由 softmax 参数化的分类分布，FIM 是稠密的，其非对角线元素 $-p_i p_j$ 耦合了所有分数。这种曲率使优化变得复杂。相比之下，作者提出了一种能产生对角 FIM的参数化方法，有效地“平坦化”了统计流形，使优化器能够沿着更直接的路径走向解。

2. 层次化二元分割

catnat函数将分类概率分布建模为一系列层次化二元决策的结果，其结构为二叉树。

• 结构：给定 $K$ 个类别，该参数化使用深度为 $H = \log_2(K)$ 的层次结构。每个未归一化分数 $s_i$ 对应树中的一个唯一节点。
• 机制：激活函数 $a(\cdot)$ 应用于每个分数以产生二元概率。类别 $k$（由二进制字符串 $\vec{b}$ 标识）的最终概率 $p_k$ 是从根节点到代表 $k$ 的叶节点路径上各概率的乘积。
• 符号：作者引入了一种符号，将二进制路径拆分为层次级别和标识符，从而允许对概率进行紧凑的数学表示。

3. 理论性质

本文证明了 catnat 参数化相对于中间分数 $\vec{s}$ 诱导了一个对角费雪信息矩阵。

• 定理 4.2：catnat 的 FIM $G_a(\vec{s})$ 是对角矩阵，其中对角线元素取决于到达节点 $i$ 的概率（$P(a_i)$）以及激活函数的导数。
• 自然激活：为了进一步简化优化景观，作者提出了一种特定的“自然激活函数”$\nu(x)$，即平滑阶跃函数。当与 catnat 结合使用时，该激活函数使 FIM 的对角线元素变为常数（仅取决于祖先概率），从而有效地消除了对局部分数值的依赖。这使得参数化与天然梯度方法的目标保持一致，而无需显式计算和求逆 FIM。

主要贡献

1. 重审 Softmax：本文提供了对 softmax 函数的信息几何批判，证明了其稠密的 FIM 在优化景观中引入了不必要的几何畸变。
2. Catnat 提案：提出了一种catnat层次化二元分割参数化方法，保证相对于输入分数具有对角 FIM。
3. 自然激活：定义了一种特定的激活函数 $\nu$，当与 catnat 配对时，可进一步简化 FIM，从对角线项中移除局部分数依赖。
4. 实现：该方法被呈现为现有代码库中 softmax 的即插即用替代品，兼容标准的训练稳定技术（例如控制变量、Gumbel-Softmax）以及各种优化器。

实验结果

作者在三个不同领域评估了 catnat 与 softmax、sparsemax 以及带 sigmoid 激活的 catnat 的表现：

1. 图结构学习（GSL）

• 设置：学习潜在图结构，其中边被建模为伯努利变量。任务涉及恢复真实的图结构和 GNN 参数。
• 结果：自然激活的 catnat 在所有指标（能量分数、点预测 MAE/MSE 以及参数估计 MAE）上始终优于 softmax、sparsemax 和基于 sigmoid 的 catnat。它在恢复真实潜在分布参数方面表现出更高的准确性。

2. 变分自编码器（VAEs）

• 设置：在 MNIST 和 Binary MNIST 数据集上的离散 VAE，具有不同数量的潜在变量（$N$）和类别（$K$）。
• 结果：在所有配置（$N \in \{10, 20, 30\}$，$K \in \{8, 16, 32\}$）中，catnat 的两个变体（sigmoid 和自然激活）均实现了比 softmax 和 sparsemax 更低的测试负对数似然值。自然激活在大多数设置中显示出轻微改进。这些结果对 Adam 和 SGD 优化器均成立，表明其优势并非特定于某种优化器。

3. 强化学习（RL）

• 设置：使用近端策略优化（PPO）在 Atari 环境（Breakout 和 Seaquest）中进行策略学习。
• 结果：带有自然激活的 catnat 参数化产生的回合回报高于标准 softmax。在 Breakout 中改进幅度适中，但在更复杂的 Seaquest 环境中改进显著，表明其在高维、序列决策任务中具有鲁棒性。

意义与主张

本文声称，用 catnat 参数化替换标准 softmax，为训练具有潜在分类变量的模型提供了一种有原则的、理论基础的改进。通过诱导对角费雪信息矩阵，catnat 减少了参数空间中的几何畸变，促进了更稳定、更高效的基于梯度的优化。

作者强调：

• 该方法与现有的方差缩减技术（例如控制变量）是互补的，可以与之结合使用。
• 该方法易于实现，并能无缝集成到现有的深度学习流程中。
• 改进具有鲁棒性，在各种架构、优化器和任务复杂度中均被观察到。

本文总结道，虽然研究聚焦于分类分布，但研究结果表明，诱导对角 FIM 的参数化通常能提升性能，这为将此方法扩展到其他分布族开辟了途径。作者将 catnat 定位为 softmax 的更优替代方案，旨在提高依赖离散潜在变量的模型的训练稳定性和最终测试性能。