一、介绍
Gumbel-Softmax 是一种技术,用于在离散选择中引入可微分的近似。这对于需要在神经网络中进行离散采样(如分类任务或生成离散数据)而不破坏梯度计算非常有用。Gumbel-Softmax 可以看作是对经典的 Softmax 函数的一种扩展,结合了 Gumbel 噪声,用于逼近离散的 one-hot 向量,同时保持梯度的可计算性。
在许多机器学习任务中,需要从一个离散的分布中采样。例如,在强化学习或生成模型中,可能需要从一组离散的动作或词汇中进行选择。然而,直接从离散分布中采样是不连续的,这意味着无法通过反向传播来更新模型参数。
举个例子:
在分类任务中,神经网络的最后一层通常是一个全连接层,接着是一个Softmax函数,将网络输出转化为概率分布。例如,对于一个有3个类别的分类任务,网络的输出可能是:
logits=[1.2,0.9,2.5]
通过Softmax函数将其转化为概率:
probs=Softmax(logits)=[0.25,0.20,0.55]
然后,通常选择概率最大的类别作为预测结果:
prediction=argmax(probs)=2
然而,离散操作(如argmax)是不可微的,这意味着无法通过反向传播来更新参数。
假设有一个简单的损失函数,它依赖于网络的输出类别。如果直接使用 argmax来选择类别,梯度无法通过这个操作传递回网络的参数:
import torch
import torch.nn.functional as F
logits = torch.tensor([1.2, 0.9, 2.5])
probs = F.softmax(logits, dim=0)
prediction = torch.argmax(probs)
# 假设有一个简单的损失函数
loss = torch.tensor(0.0)
if prediction == 2:
loss = torch.tensor(1.0)
# 反向传播
loss.backward() # 这里将会失败,因为 torch.argmax 是不可微的
二、Gumbel-Softmax 的作用
Gumbel-Softmax 提供了一种可微的近似,使可以在保持反向传播可行的情况下进行离散选择。通过Gumbel-Softmax,可以从离散分布中进行采样,并且保持梯度的传递。
三、Gumbel-Softmax 的工作原理
-
添加Gumbel噪声:首先在logits上添加Gumbel噪声,使其具有随机性。
-
应用Softmax:然后应用Softmax函数,使得输出接近one-hot向量,但仍然是连续的,从而可以进行梯度计算。
import torch
import torch.nn.functional as F
def sample_gumbel(shape, eps=1e-20):
U = torch.rand(shape)
return -torch.log(-torch.log(U + eps) + eps)
def gumbel_softmax_sample(logits, temperature):
gumbel_noise = sample_gumbel(logits.size())
y = logits + gumbel_noise
return F.softmax(y / temperature, dim=-1)
def gumbel_softmax(logits, temperature, hard=False):
y = gumbel_softmax_sample(logits, temperature)
if hard:
shape = y.size()
_, max_idx = y.max(dim=-1, keepdim=True)
y_hard = torch.zeros_like(y).scatter_(-1, max_idx, 1.0)
y = (y_hard - y).detach() + y
return y
logits = torch.tensor([1.2, 0.9, 2.5])
temperature = 0.5
samples = gumbel_softmax(logits, temperature, hard=True)
print(samples)
# 我们有一个简单的损失函数
target = torch.tensor([0, 0, 1]) # 假设真实类别是第三类
loss = F.cross_entropy(samples.unsqueeze(0), target.unsqueeze(0).argmax(dim=-1))
# 反向传播
loss.backward() # 这里是可行的,因为Gumbel-Softmax是可微的
四、公式推导
假设有一个logits向量 z(通常是神经网络的输出),其第 i 个元素为 zi。
添加Gumbel噪声: 对每个logit zi 添加一个从Gumbel分布中采样的噪声 gi。Gumbel噪声 gi 的公式为:
其中 U 是从均匀分布 U(0,1) 中采样的随机变量
应用Softmax: 对添加了Gumbel噪声的logits进行缩放,并应用Softmax函数:
其中,τ 是温度参数,控制输出的平滑程度。当τ 趋近于0时,输出接近于one-hot向量;当τ 趋近于无穷大时,输出接近均匀分布。
对于 Gumbel-Softmax,梯度的计算与 Softmax 类似,只不过在 logits 中添加了 Gumbel 噪声。具体来说,对于 Gumbel-Softmax 的输出 yi 和输入 logits zi,其梯度计算如下:
本质上还是应用 Softmax 的梯度公式,但在计算过程中包含了 Gumbel 噪声。