NO.35.tip：随机优化

背景

我们将讨论随机目标的优化，其形式如下：

$L (θ) = E_{q_{θ} (z)} [L (θ, z)] (1)$

其中， $θ$ 是我们正在优化的参数， $z$ 是随机变量，例如外部噪声。

随机梯度下降

假设我们有一种计算目标函数梯度的无偏估计 $g_{t}$ 的方法，即

$E [g_{t}] = \nabla_{θ} L (θ) ∣_{θ_{t}} (2)$

于是，我们可以在梯度下降过程中使用该估计：

$θ_{t + 1} = θ_{t} - η_{t} g_{t} (3)$

其中， $η_{t}$ 是学习率（learning rate）或步长（step size）。这被称为随机梯度下降。

源与流

选择步长

在使用随机梯度下降时，为了实现收敛，在选择学习率时需要特别小心。我们可以使用学习率调度（learning rate schedule），而不是选择单一的恒定学习率。在学习率调度中，我们会随着时间的推移调整步长。理论上，随机梯度下降实现收敛的一个充分条件是，如果学习率调度满足 Robbins-Monro 条件：

$η_{t} \to 0, \frac{\sum _{t = 1}^{\infty} η _{t}^{2}}{\sum _{t = 1}^{\infty} η _{t}} \to 0 (4)$

下面列出了一些学习率调度的常见示例：

$η_{t} = η_{i} 如果 t_{i} \leq t \leq t_{i + 1} 分段常数 (5)$

$η_{t} = η_{0} e^{- λ t} 指数衰减 (6)$

$η_{t} = η_{0} (βt + 1)^{- α} 多项式衰减 (7)$

在分段常数调度中， $t_{i}$ 是一组时间点，在这些时间点上，我们将学习率调整到指定值。例如，我们可以设置 $η_{t} = η_{0} γ^{i}$ ，这将使我们通过的每个阈值（或里程碑）的初始学习率降低 $γ$ 倍。可以看出 $η_{0} = 1$ 和 $γ = 0.9$ 的情况。这被称为阶跃衰减（step decay）。有时，通过估计训练或验证损失何时稳定下来，自适应地计算阈值时间，这被称为高原衰减（reduce on-plateau）。指数衰减通常过快。常见的选择是多项式衰减，当 $α = 0.5, β = 1$ 时，这对应于平方根调度（square-root schedule） $η_{t} = η_{0} \frac{1}{t + 1}$ 。

方差缩减

由于随机梯度下降依赖于梯度的随机估计，因此可能收敛较慢。对于在每个步骤中生成的参数估计，研究人员已经提出了各种方法来减少这些参数估计的方差，这可以加速收敛。

预处理随机梯度下降

在许多情况下，梯度大小沿着每个维度可能存在较大差异，对应于损失函数的表面沿着某些方向是陡峭的，而沿着其他方向是平缓的，类似于谷底。在这种情况下，可以通过条件矩阵（conditioning matrix） $C_{t}$ 对梯度向量进行缩放来获得更快的收敛，如下所示：

$θ_{t + 1} = θ_{t} - η_{t} C_{t} g_{t} (8)$

这被称为预处理随机梯度下降。

用于优化"有限和"目标的随机梯度下降

在最简单的情况下，用于计算期望值的分布 $q_{θ} (z)$ 不取决于被优化的参数 $θ$ 。在这种情况下，我们可以将梯度移动到期望算子内部，然后对 $z$ 使用蒙特卡罗采样来近似梯度：

$\nabla_{θ} L (θ) = \nabla_{θ} E_{q (z)} [L (θ, z)] = E_{q (z)} [\nabla_{θ} L (θ, z)] \approx \frac{1}{S} s = 1 \sum S \nabla_{θ} L (θ, z_{s}) (9)$

例如，考虑经验风险最小化（Empirical Risk Minimization, ERM）问题，该问题需要最小化以下目标：

$L (θ) = \frac{1}{N} n = 1 \sum N L (θ, z_{n}) = \frac{1}{N} n = 1 \sum N ℓ (y_{n}, f (x_{n}; θ)) (10)$

其中， $z_{n} = (x_{n}, y_{n})$ 是第 $n$ 个标记样例， $f$ 是预测函数。这种目标被称为"有限和"（finite sum objective）。我们可以将其记作经验分布 $p_{D} (x, y)$ 的预期损失：

$L (θ) = E_{p_{D} (z)} [L (θ, z)] (11)$

由于期望值取决于数据，而不是参数，因此我们可以在每次迭代时使用来自完整数据集的 $B = ∣ B ∣$ 个数据点的迷你批次来近似梯度：

$g_{t} = \nabla L (θ_{t}) = \frac{1}{B} n \in B \sum \nabla ℓ (y_{n}, f (x_{n}; θ)) (12)$

然后可以将这些噪声梯度传递给随机梯度下降。当数据集很大时，这种方法比全批次（full batch）梯度下降要快得多，因为在更新模型之前，该方法不需要评估所有 $N$ 个样例的损失。

用于优化分布参数的随机梯度下降

现在假设随机性取决于我们正在优化的参数。例如， $z$ 可以是从随机策略 $q_{θ}$ 中采样的行为，或者 $z$ 可以是在随机变分推理中从推理网络 $q_{θ}$ 采样的潜在变量。在这种情况下，梯度的计算公式为：

$\nabla_{θ} E_{q_{θ} (z)} [L (θ, z)] = \nabla_{θ} \int L (θ, z) q_{θ} (z) d z (13)$

$= \int [\nabla_{θ} L (θ, z)] q_{θ} (z) d z + \int L (θ, z) [\nabla_{θ} q_{θ} (z)] d z (14)$

第一项可以通过蒙特卡罗采样来近似：

$\int [\nabla_{θ} L (θ, z)] q_{θ} (z) d z \approx \frac{1}{S} s = 1 \sum S \nabla_{θ} L (θ, z_{s}) (15)$

其中， $z_{s} \sim q_{θ}$ 。请注意，如果 $L ()$ 与 $θ$ 无关，那么该项将消除。

现在考虑第二项，该项取分布本身的梯度：

$I ≜ \int L (θ, z) [\nabla_{θ} q_{θ} (z)] d z (16)$

我们不能再使用基本的蒙特卡罗采样来近似这个积分。然而，还有各种其他的方法来近似。

得分函数估计器

近似式 (16) 的最简单方法是利用对数导数技巧（log derivative trick），该技巧是以下恒等式：

$\nabla_{θ} q_{θ} (z) = q_{θ} (z) \nabla_{θ} lo g q_{θ} (z) (17)$

由此，我们可以将式 (16) 改写如下：

$I = \int L (θ, z) [q_{θ} (z) \nabla_{θ} lo g q_{θ} (z)] d z = E_{q_{θ} (z)} [L (θ, z) \nabla_{θ} lo g q_{θ} (z)] (18)$

这被称为得分函数估计器（Score Function Estimator, SFE）。（术语"得分函数"指的是对数概率分布的梯度）得分函数估计器也被称为似然比梯度估计器（likelihood ratio gradient estimator），或 REINFORCE 估计器。我们现在可以很容易地使用蒙特卡罗进行近似：

$I \approx \frac{1}{S} s = 1 \sum S L (θ, z_{s}) \nabla_{θ} lo g q_{θ} (z_{s})$

其中， $z_{s} \sim q_{θ}$ 。我们只要求采样分布是可微的，而不要求目标函数 $L (θ, z)$ 本身是可微的。因此，该方法可用于黑盒随机优化问题，例如变分优化、黑盒变分推理、强化学习等。

控制变量法

利用得分函数进行估计可能具有高方差。减少这种情况的一种方法是使用控制变量法，其中我们将 $L (θ, z)$ 替换为：

$L (θ, z) = L (θ, z) - c (b (θ, z) - E [b (θ, z)]) (19)$

在上式中， $b (θ, z)$ 是与 $L (θ, z)$ 相关的基线函数（baseline function），并且 $c < 0$ 是一个系数。由于 $E [L (θ, z)] = E [L (θ, z)]$ ，我们可以使用 $L$ 来计算 $L$ 的无偏梯度估计。该方法的优越性在于，这种新的估计可以导致较低的方差。

Rao-Blackwellization

假设 $q_{θ} (z)$ 是一个离散分布。在这种情况下，我们的目标变为 $L (θ) = \sum_{z} L (θ, z) q_{θ} (z)$ 。现在，我们可以使用 $\nabla_{θ} L (θ) = \sum_{z} L (θ, z) \nabla_{θ} q_{θ} (z)$ 来简单地计算梯度。当然，如果 $z$ 可以取指数级的多个值（例如，我们在字符串空间上进行优化），那么这个表达式是难以处理的。然而，假设我们可以将这个求和划分为两个集合，一个是高概率值的小集合 $S_{1}$ ，另一个是所有其他值的大集合 $S_{2}$ 。于是，我们可以在 $S_{1}$ 上枚举，并在 $S_{2}$ 上使用得分函数估计器：

$\nabla_{θ} L (θ) = z \in S_{1} \sum L (θ, z) \nabla_{θ} q_{θ} (z) + E_{q_{θ} (z ∣ z \in S_{2})} [L (θ, z) \nabla_{θ} lo g q_{θ} (z)] (20)$

为了计算第二项的期望值，我们可以使用应用于来自 $q_{θ} (z)$ 的样本的拒绝采样。该程序是 Rao-Blackwellization 的一种形式，与标准得分函数估计器相比，这种方法减少了方差。

重新参数化的技巧

即使在使用控制变量法时，得分函数估计器也可能具有高方差。在本节中，我们将推导出一个低方差估计器，如果 $L (θ, z)$ 相对于 $z$ 是可微的，则可以应用该估计器。另外，我们还要求，可以通过首先从独立于 $θ$ 的某个噪声分布 $q_{0}$ 中采样，然后使用确定性和可微函数 $z = g (θ, ϵ)$ 将其转换为 $z$ ，以从 $q_{θ} (z)$ 计算样本。例如，我们可以采样 $ϵ \sim N (0, 1)$ 并计算，而不是采样 $z \sim N (μ, σ^{2})$ ：

$z = g (θ, ϵ) = μ + σ ϵ (21)$

其中， $θ = (μ, σ)$ 。这允许我们将随机目标改写如下：

$L (θ) = E_{q_{θ} (z)} [L (θ, z)] = E_{q_{0} (ϵ)} [L (θ, g (θ, ϵ))] (22)$

由于 $q_{0} (ϵ)$ 与 $θ$ 无关，我们可以将梯度算子移动到期望内，从而可以使用蒙特卡罗近似：

$\nabla_{θ} L (θ) = E_{q_{0} (ϵ)} [\nabla_{θ} L (θ, g (θ, ϵ))] \approx \frac{1}{S} s = 1 \sum S \nabla_{θ} L (θ, g (θ, ϵ_{s})) (23)$

其中， $ϵ_{s} \sim q_{0}$ 。这被称为重新参数化梯度（reparameterization gradient）或路径导数（path-wise derivative），并被广泛用于变分推理。

示例

作为一个简单的例子，假设我们定义了某种任意函数，例如 $L (z) = z^{2} - 3 z$ ，然后将其期望值定义为 $L (θ) = E_{N (z ∣ μ, v)} [L (z)]$ ，其中 $θ = (μ, v)$ ， $v = σ^{2}$ 。假设我们要计算：

$\nabla_{θ} L (θ) = [\frac{\partial}{\partial μ} E [L (z)], \frac{\partial}{\partial v} E [L (z)]] (24)$

由于高斯分布是可重新参数化的，因此我们可以对 $z \sim N (z ∣ μ, v)$ 进行采样，然后使用自动微分来计算这些梯度项，之后进行平均。

然而，在高斯分布的特殊情况下，我们也可以直接计算梯度向量。有Bonnet 定理，该定理表明：

$\frac{\partial}{\partial μ} E [L (z)] = E [\frac{\partial}{\partial z} L (z)] (25)$

类似地，Price 定理表明：

$\frac{\partial}{\partial v} E [L (z)] = 0.5 E [\frac{\partial ^{2}}{\partial z ^{2}} L (z)] (26)$

我们可以证明这两种方法在数值上是等价的，正如理论所表明的结论。

全导数

为了计算式 (23) 中期望内的梯度项，我们需要使用全导数（total derivative），因为函数 $L$ 通过噪声样本 $z$ 直接依赖于 $θ$ 。回想一下，对于形式为 $L (θ_{1}, \dots, θ_{d_{ψ}}, z_{1} (θ), \dots, z_{d_{z}} (θ))$ 的函数，相对于 $θ_{i}$ 的全导数可以通过链式法则进行计算：

$\frac{\partial L}{\partial θ _{i}}^{TD} = \frac{\partial L}{\partial θ _{i}} + j \sum \frac{\partial L}{\partial z _{j}} \frac{\partial z _{j}}{\partial θ _{i}} (27)$

因此：

$\nabla_{θ} L (θ, z)^{TD} = \nabla_{z} L (θ, z) J + \nabla_{θ} L (θ, z) (28)$

其中， $J = \frac{\partial z ^{T}}{\partial θ}$ 是噪声转换的大小为 $d_{z} \times d_{ψ}$ 的雅可比矩阵。

$J = \frac{\partial z _{1}}{\partial θ _{1}} ⋮ \frac{\partial z _{d_{z}}}{\partial θ _{1}} \dots \dots \frac{\partial z _{1}}{\partial θ _{d_{ψ}}} ⋮ \frac{\partial z _{d_{z}}}{\partial θ _{d_{ψ}}} (29)$

"稳着地"估计器

我们考虑变分推理中出现的特殊情况。证据下界目标（对于单个潜在样本 $z$ ）具有以下形式：

$L (θ, z) = lo g p (z, x) - lo g q (z ∣ θ) (30)$

其中， $θ$ 是变分后验的参数。则梯度变为：

$\nabla_{θ} L (θ, z) = \nabla_{θ} [lo g p (z, x) - lo g q (z ∣ θ)] (31)$

$= \nabla_{z} [lo g p (z, x) - lo g q (z ∣ θ)] J - \nabla_{θ} lo g q (z ∣ θ) (32)$

第一项是 $θ$ 通过生成的样本 $z$ 对目标的间接影响。第二项是 $θ$ 对目标的直接影响。第二项在期望中为零，因为该项是评分函数，但对于有限数量的样本，该项可能是非零的，即使 $q (z ∣ θ) = p (z ∣ x)$ 是真正的后验。我们建议去掉第二项以创建一个较低的方差估计器。这可以通过使用 $lo g q (z ∣ θ^{'})$ 来实现，其中 $θ^{'}$ 是 $θ$ 的"断开"副本，不影响梯度的值。在伪代码中，该方法如下所示：

$ϵ \sim q_{0} (ϵ) (33)$

$z = g (ϵ, θ) (34)$

$θ^{'} = stop-gradient (θ) (35)$

$g = \nabla_{θ} [lo g p (z, x) - lo g q (z ∣ θ^{'})] (36)$

我们称之为"稳着地"（Sticking The Landing, STL）估计器。请注意，"稳着地"估计器并不总是比没有停止梯度项的"标准"估计器好。我们建议使用估计量的加权组合，其中对权重进行优化，以减少固定计算量的方差。

Gumbel softmax技巧

在处理离散变量时，我们不能使用重新参数化技巧。然而，我们通常可以将离散变量放松为连续变量，这样就可以使用重新参数化技巧，如下所述。

考虑一个具有 $K$ 个比特的独热向量 $d$ ，因此 $d_{k} \in {0, 1}$ ， $\sum_{k = 1}^{K} d_{k} = 1$ 。这可以用来表示 $K$ 元分类变量 $d$ 。设 $P (d) = Cat (d ∣ π)$ ，其中 $π_{k} = P (d_{k} = 1)$ ，因此 $0 \leq π_{k} \leq 1$ 。或者，我们可以使用 $(α_{1}, \dots, α_{K})$ 来参数化分布，其中 $π_{k} = α_{k} / (\sum_{k^{'} = 1}^{K} α_{k^{'}})$ 。我们将使用 $d \sim Cat (d ∣ α)$ 来表示。

我们可以通过以下计算来从该分布中采样独热向量 $d$ ：

$d = onehot (ar g k max [ϵ_{k} + lo g α_{k}]) (37)$

其中， $ϵ_{k} \sim Gumbel (0, 1)$ 从 Gumbel 分布中采样。我们可以通过首先采样 $u_{k} \sim Unif (0, 1)$ ，然后计算 $ϵ_{k} = - lo g (- lo g (u_{k}))$ 来采样这样的样本。这被称为 Gumbel 最大技巧（Gumbel-max trick），它为我们提供了分类分布的可重新参数化表示。

遗憾的是， $ar g max$ 的导数在任何地方都是 0，除了在从一个标签到另一个标签的转换边界处，该转换边界处的导数是未定义的。然而，假设我们使用 $softmax$ 替换 $ar g max$ ，并使用连续松弛 $x \in Δ^{K - 1}$ 替换离散的独热向量 $d$ ，其中 $Δ^{K - 1} = {x \in R^{K} : x_{k} \in [0, 1], \sum_{k = 1}^{K} x_{k} = 1}$ 是 $K$ 维单纯形，那么可以记作：

$x_{k} = \frac{exp (( lo g α _{k} + ϵ _{k} ) / τ )}{\sum _{k^{'} = 1}^{K} exp (( lo g α _{k^{'}} + ϵ _{k^{'}} ) / τ )} (38)$

其中， $τ < 0$ 是一个温度参数。这被称为 Gumbel softmax 分布或具体分布（concrete distribution）。当 $τ \to 0$ 时，该分布平滑地接近离散分布。

我们现在可以使用 $f (d)$ 代替 $f (x)$ ，这允许我们采用重新参数化的梯度（相对于 $x$ ）。

随机计算图

我们可以将包含确定性和随机分量的任意函数表示为随机计算图（stochastic computation graph）。然后，我们可以推广自动微分算法，利用得分函数估计和重新参数化来计算复杂嵌套函数的蒙特卡罗梯度。

直通式估计器

我们将讨论如何对信号的量化版本的梯度进行近似。例如，假设我们有以下阈值函数，该函数对其输出进行二值化：

$f (x) = {10 如果 x < 0 如果 x \leq 0 (39)$

上式并没有一个明确定义的梯度。然而，我们提出的直通式估计器（straight-through estimator）作为近似。其基本思想是，在计算向后传递时，将 $g (x) = f^{'} (x)$ （其中， $f^{'} (x)$ 是 $f$ 相对于输入的导数）替换为 $g (x) = x$ 。在实践中，我们有时使用 hard tanh 函数代替 $g (x) = x$ ，其定义如下：

$HardTanh (x) = ⎩ ⎨ ⎧ x 1 - 1 如果 - 1 \leq x \leq 1 如果 x < 1 如果 x lt; - 1 (40)$

这样可以确保反向传播的梯度不会变得太大。

概率统计100问

NO.35.tip： 随机优化

背景

源与流

NO.35.tip：随机优化