原理推导
1.互信息
2.信息瓶颈
也就是我们假设从输出到特征空间再到输出是一个马尔科夫过程
现在我们的目标就是让Z充分包含Y的信息，同时为了压缩效果是X和Z的互信息尽可能少
也就是：
引入拉格朗日乘子，即最大化目标：
将公式进一步展开：
下面介绍变分信息瓶颈：
现在我们对X,Z,Y的马尔科夫链进一步假设：
这其实是符合神经网络模型的：
因此，我们有：
我们可以进一步得到：
因为 p(y|z) 无法直接计算，假设 q(y|z) 是 p(y|z) 的变分近似，也就是我们模型中的encoder
利用KL散度非负的特性，可得：
我们对I(Z,Y)化简：
将 p(y,z) 写成 p(y,z)=p(x)p(y|x)p(z|x) ，可以得到新的下界：
现在我们考虑化简I(Z,X)
同理，让r(z)作为p(z)的变分近似：
因此，可得：
现在我们代入上述不等式对信息瓶颈的公式进行化简：
在实际计算中，由于我们只能对样本采样来拟合分布，因此：
Xn,Yn也就是数据集中的输入和输出的样本
进一步代入公式：
我们将输入X，通过encoder得到z的概率分布，然后进行采样，即：
考虑到“采样”这个操作是随机的，我们无法直接对它求导，于是进行参数重整化
最终，我们最小化：
其中：
下面展示一下怎么从代码层面进行实现（主要讲一下核心的函数）：
KL散度函数
# 定义函数，接收两个分布的参数
def KL_between_normals(q_distr, p_distr):
# 1. 解包输入参数
# mu_q, sigma_q 的形状通常是 [batch_size, k]，k是潜变量维度
# sigma_q 是标准差 (standard deviation)，而不是方差 (variance)
mu_q, sigma_q = q_distr
mu_p, sigma_p = p_distr
# 2. 获取潜变量的维度 k
k = mu_q.size(1)
# 3. 计算均值差的平方项：(μ_p - μ_q)²
# 这一步是逐元素计算的
mu_diff = mu_p - mu_q
mu_diff_sq = torch.mul(mu_diff, mu_diff) # 等价于 mu_diff ** 2
# 4. 计算两个分布的对数行列式 (log-determinant)
# 对于对角协方差矩阵 Σ = diag(σ₁², σ₂², ...)，其行列式 det(Σ) = Π σᵢ²
# 对数行列式 log(det(Σ)) = log(Π σᵢ²) = Σ log(σᵢ²) = Σ 2*log(σᵢ)
# torch.clamp 是为了数值稳定性，防止 sigma 变为 0 导致 log(0) = -inf
logdet_sigma_q = torch.sum(2 * torch.log(torch.clamp(sigma_q, min=1e-8)), dim=1)
logdet_sigma_p = torch.sum(2 * torch.log(torch.clamp(sigma_p, min=1e-8)), dim=1)
# 5. 计算公式中的两个主要分数项
# fs 对应于公式中的 Σ [ (σ_q²/σ_p²) + ((μ_p - μ_q)²/σ_p²) ]
# torch.div 是逐元素除法
# torch.sum(..., dim=1) 是沿着维度k进行求和，得到每个样本的总和
# (sigma_q ** 2) 是方差 σ²
term1 = torch.sum(torch.div(sigma_q ** 2, sigma_p ** 2), dim=1) # 对应 tr(Σ_p⁻¹ * Σ_q)
term2 = torch.sum(torch.div(mu_diff_sq, sigma_p ** 2), dim=1) # 对应 (μ_p - μ_q)ᵀ Σ_p⁻¹ (μ_p - μ_q)
fs = term1 + term2
# 6. 组合所有部分
# 这里计算的是 2 * KL(q || p)
# 2*KL = Σ [ σ_q²/σ_p² + (μ_p-μ_q)²/σ_p² - 1 - log(σ_q²) + log(σ_p²) ]
# = (fs) - Σ(1) + Σ(log(σ_p²)) - Σ(log(σ_q²))
# = fs - k + logdet_sigma_p - logdet_sigma_q
two_kl = fs - k + logdet_sigma_p - logdet_sigma_q
# 7. 返回最终的 KL 散度值
# 将结果除以2，得到最终的KL散度。
# 输出的张量形状为 [batch_size]，每个元素是对应样本的KL散度值
return two_kl * 0.5
解释：
KL数学解释：
VIB模型：
class VIB(nn.Module):
def __init__(self, X_dim, y_dim, dimZ=256, beta=1e-3, num_samples=10):
super().__init__()
# 1. 保存超参数
self.beta = beta # VIB损失函数中的权重 β
self.dimZ = dimZ # 潜变量 Z 的维度
self.num_samples = num_samples # 蒙特卡洛采样次数，用于估算期望
# 2. 定义编码器
# 这是一个简单的多层感知机 (MLP)
self.encoder = nn.Sequential(nn.Linear(in_features=X_dim, out_features=1024),
nn.ReLU(),
nn.Linear(in_features=1024, out_features=1024),
nn.ReLU(),
# 关键: 输出维度是 2 * dimZ
# 一半用于均值(mu)，一半用于标准差(sigma)
nn.Linear(in_features=1024, out_features=2 * self.dimZ))
# 3. 定义解码器
# 论文中使用了简单的逻辑回归，这里就是一个线性层
# 它接收潜变量 Z (dimZ)，输出预测 Y 的 logits (y_dim)
self.decoder_logits = nn.Linear(in_features=self.dimZ, out_features=y_dim)
def gaussian_noise(self, num_samples, K):
# works with integers as well as tuples
return torch.normal(torch.zeros(*num_samples, K), torch.ones(*num_samples, K)).cuda()
def sample_prior_Z(self, num_samples):
return self.gaussian_noise(num_samples=num_samples, K=self.dimZ)
def encoder_result(self, batch):
encoder_output = self.encoder(batch)
mu = encoder_output[:, :self.dimZ]
sigma = torch.nn.functional.softplus(encoder_output[:, self.dimZ:])
return mu, sigma
def sample_encoder_Z(self, num_samples, batch):
batch_size = batch.size()[0]
mu, sigma = self.encoder_result(batch)
return mu + sigma * self.gaussian_noise(num_samples=(num_samples, batch_size), K=self.dimZ)
def forward(self, batch_x):
batch_size = batch_x.size()[0]
# sample from encoder
encoder_Z_distr = self.encoder_result(batch_x)
to_decoder = self.sample_encoder_Z(num_samples=self.num_samples, batch=batch_x)
decoder_logits_mean = torch.mean(self.decoder_logits(to_decoder), dim=0)
return decoder_logits_mean
def batch_loss(self, num_samples, batch_x, batch_y):
batch_size = batch_x.size()[0]
# --- 第一部分：计算 KL 散度项 I(Z;X) ---
# 1. 定义先验分布 N(0, I)
prior_Z_distr = torch.zeros(batch_size, self.dimZ).cuda(), torch.ones(batch_size, self.dimZ).cuda()
# 2. 获取编码器产生的后验分布 p(z|x)
encoder_Z_distr = self.encoder_result(batch_x)
# 3. 计算 KL(p(z|x) || p(z)) 并对批次取平均
I_ZX_bound = torch.mean(KL_between_normals(encoder_Z_distr, prior_Z_distr))
# --- 第二部分：计算交叉熵项 -I(Z;Y) ---
# 1. 使用重参数化技巧从 p(z|x) 中采样Z
# to_decoder 形状: (num_samples, batch_size, dimZ)
to_decoder = self.sample_encoder_Z(num_samples=self.num_samples, batch=batch_x)
# 2. 将Z送入解码器得到预测
# decoder_logits 形状: (num_samples, batch_size, y_dim)
decoder_logits = self.decoder_logits(to_decoder)
# 3. 准备计算交叉熵损失
# nn.CrossEntropyLoss期望的输入形状是 (N, C, ...)，其中N是批大小，C是类别数
# 当前形状是 (num_samples, batch_size, y_dim)
# 我们把它变成 (batch_size, y_dim, num_samples) 来匹配PyTorch的要求
decoder_logits = decoder_logits.permute(1, 2, 0)
# 4. 计算每个样本、每个采样的交叉熵损失
# reduce=False: 不对损失求和或求平均，返回每个元素的损失
loss = nn.CrossEntropyLoss(reduce=False)
# batch_y 形状是 (batch_size)，需要扩展成 (batch_size, num_samples)
# 以便和 (batch_size, y_dim, num_samples) 的预测logits对应
cross_entropy_loss = loss(decoder_logits, batch_y[:, None].expand(-1, num_samples))
# 5. 通过对采样维度求平均来估算期望 E[log q(y|z)]
cross_entropy_loss_montecarlo = torch.mean(cross_entropy_loss, dim=-1)
# 6. 对批次求平均，得到最终的交叉熵损失项
minusI_ZY_bound = torch.mean(cross_entropy_loss_montecarlo, dim=0)
# --- 合并两部分，得到最终的总损失 ---
# J_IB = -I(Z;Y) + β * I(Z;X)
total_loss = minusI_ZY_bound + self.beta * I_ZX_bound
# 返回总损失和两个分量，方便监控训练过程
return total_loss, -minusI_ZY_bound, I_ZX_bound
参考代码链接：https://github.com/makezur/VIB_pytorch/blob/master/VIB.ipynb
扫一扫
1112
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
