监督式VAE分类器的损失权重设置及最优权重求解方法

我在音频分类领域开展研究，近期尝试使用Supervised VAE Classifier，所用模型架构如下：

import torch
import torch.nn as nn
import torch.nn.functional as F

class VAE(nn.Module):
    def __init__(self, input_shape, latent_size):
        super(VAE, self).__init__()
        # Encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(input_shape[0], 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(),
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(),
            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU()
        )
        self.fc_mu = nn.Linear(512*8*1, latent_size)
        self.fc_logvar = nn.Linear(512*8*1, latent_size)

        # Decoder
        self.decoder_input = nn.Linear(latent_size, 512*8*1)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(),
            nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(),
            nn.ConvTranspose2d(64, input_shape[0], kernel_size=3, stride=2, padding=1, output_padding=(1, 0)),
            nn.Sigmoid()
        )

        # Classifier
        self.clf = nn.Sequential(
            nn.Linear(latent_size, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(),
            nn.Dropout(0.25),

            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.LeakyReLU(),

            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.LeakyReLU(),

            nn.Linear(128, 7),
        )

    def encode(self, x):
        x = self.encoder(x)
        x = x.view(-1, 512*8*1)
        mu = self.fc_mu(x)
        logvar = self.fc_logvar(x)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std

    def decode(self, z):
        x = self.decoder_input(z)
        x = x.view(-1, 512, 8, 1)
        x = self.decoder(x)
        return x

    def forward(self, x):
        mu, logvar= self.encode(x)
        z = self.reparameterize(mu, logvar)
        reconstruction = self.decode(z)
        clf = self.clf(z)
        return reconstruction, mu, logvar, clf

常规VAE训练采用BCE损失与KL散度之和，但在引入交叉熵损失的监督式VAE分类器中，如何权衡这三类损失是关键问题。我看到部分研究采用手动设置权重的方式，示例代码如下：

def vae_loss(recon_x, x, mu, logvar, clf, target):
    input_size = x.size(1) * x.size(2) * x.size(3)
    # BCE Loss
    BCE = F.binary_cross_entropy(recon_x.view(-1, input_size), x.view(-1, input_size), reduction='mean') #sum
    # Kullback-Leibler Divergence
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    # CrossEntropy Loss
    clf_loss = F.cross_entropy(clf, target)
    return 0.001*(BCE + 3*KLD)+ clf_loss

针对「是否存在求解最优损失权重的可行方案」，以下是几种实用思路：

一、手动调优的经验策略 #

• 分步训练：先单独训练VAE模块，让重构损失（BCE）和KL散度稳定收敛，再引入分类损失，此时VAE部分的权重可以从较小值开始逐步调整
• 统一损失量级：注意示例中BCE用了mean而KLD用了sum，先将所有损失的计算方式统一（比如都用mean），避免某类损失因量级过大/过小被忽略
• 固定基准权重：先将分类损失权重设为1，调整VAE部分的权重，观察验证集的分类精度与音频重构质量（比如用PSNR、SSIM指标），找到两者的平衡点

二、自适应动态权重 #

• 梯度平衡：跟踪每个损失的梯度范数，动态调整权重，让各损失对总梯度的贡献相近，避免某一任务主导训练
• 损失均值缩放：用指数移动平均计算每个损失的运行均值，将权重设为均值的倒数，让各损失的量级始终保持在同一水平
• 可学习权重：将损失权重作为可训练参数加入模型，用sigmoid或softplus激活确保权重为正，通过反向传播自动优化权重，但需注意添加正则项防止权重过大

三、自动化搜索方案 #

• 网格搜索：在合理的参数范围内（比如VAE总损失权重取[0.0001, 0.001, 0.01, 0.1]），遍历所有组合，用验证集的分类精度+重构质量作为评价指标选择最优权重
• 贝叶斯优化：用优化工具自动搜索最优权重组合，相比网格搜索更高效，适合多参数的权重调优场景

四、损失标准化预处理 #

对每个损失项进行标准化，比如除以该损失在初始训练几步的均值，让所有损失的初始量级一致，之后可以用相等的权重训练，后期再根据验证结果微调

内容的提问来源于stack exchange，提问作者randomboar