为变分自编码器（VAE）指定输入特征赋予更高权重以增强数据多样性

问题背景 #

我用变分自编码器（VAE）做数据增强，模型能正常运行，但生成的数据里第二列特征的变化太小，当前生成结果示例：

100  16 2.6
105  16.6 2.7
110  16.7 2.8

输入数据中第二列的范围是1-100，但生成时没能覆盖这个范围，输入数据示例：

100 16 2.5
110 20 3.5
120 30 3.7
130 40 4 
200 80 7 
.....

我希望第二列能有更多变化，覆盖1-100的完整范围，期望生成结果类似：

100 20  3
105 50  3.5
110 80  6
....

我尝试了自定义损失函数，但效果不好，代码如下：

class customLoss(nn.Module):
    def __init__(self):
        super(customLoss, self).__init__()
        self.mse_loss = nn.MSELoss(reduction="sum")
    
    def forward(self, x_recon, x, mu, logvar):
        loss_MSE = self.mse_loss(x_recon, x)
        #   Kullback-Leibler (KL) divergence
        loss_KLD = -0.5 * torch.sum(1. + logvar - mu.pow(2) - logvar.exp())
        # Variance penalty for feature 2
        variance_penalty = torch.mean(logvar[:, 2])  # Average logvar for feature 2

        loss_original = loss_MSE + loss_KLD - variance_penalty * 2.
        return loss_original

改进方案 #

1. 修正自定义损失函数并加权特征 #

原损失函数存在索引错误（把第三列当成了第二列），同时可以给第二列的重建损失和方差惩罚设置针对性权重：

class customLoss(nn.Module):
    def __init__(self):
        super(customLoss, self).__init__()
        self.mse_loss = nn.MSELoss(reduction="none")  # 保留逐元素损失，方便按列加权
    
    def forward(self, x_recon, x, mu, logvar):
        # 给第二列（索引1，Python从0开始计数）设置更高的重建权重
        feature_weights = torch.tensor([1.0, 5.0, 1.0], device=x.device)
        mse_per_feature = self.mse_loss(x_recon, x) * feature_weights
        loss_MSE = torch.sum(mse_per_feature)
        
        # KL散度保持原有计算逻辑
        loss_KLD = -0.5 * torch.sum(1. + logvar - mu.pow(2) - logvar.exp())
        
        # 修正索引：针对第二列的logvar做惩罚，鼓励其增大方差
        variance_penalty = torch.mean(logvar[:, 1])
        # 调整惩罚系数，越大越能推动第二列生成更多变化
        loss_total = loss_MSE + loss_KLD - variance_penalty * 5.0
        
        return loss_total

• 关键修正：把原代码的logvar[:,2]改为logvar[:,1]，对应真实的第二列。
• 权重系数（如5.0）可根据生成效果逐步调整，数值越大，模型越偏向于让第二列产生波动。

2. 预处理阶段放大第二列的缩放范围 #

在数据归一化时，给第二列设置更大的缩放区间，让模型训练时更容易捕捉该特征的波动：

# 假设输入数据为numpy数组data
from sklearn.preprocessing import MinMaxScaler

scalers = []
scaled_data = []
for col_idx in range(data.shape[1]):
    scaler = MinMaxScaler()
    if col_idx == 1:
        # 第二列缩放至[-2,2]区间，远大于其他列的[0,1]，扩大波动空间
        scaled_col = scaler.fit_transform(data[:, col_idx].reshape(-1,1)) * 4 - 2
    else:
        scaled_col = scaler.fit_transform(data[:, col_idx].reshape(-1,1))
    scaled_data.append(scaled_col)
    scalers.append(scaler)

scaled_data = np.concatenate(scaled_data, axis=1)

生成数据后，再用对应的scaler逆变换回原范围即可。

3. 解码器输出层针对性初始化 #

在解码器的输出层，给第二列对应的神经元设置更大的初始权重，让该特征更容易产生大的数值波动：

class Decoder(nn.Module):
    def __init__(self, latent_dim, input_dim):
        super().__init__()
        self.fc1 = nn.Linear(latent_dim, 64)
        self.fc2 = nn.Linear(64, input_dim)
        
        # 给第二列对应的输出权重乘3，放大初始波动能力
        with torch.no_grad():
            self.fc2.weight[1] *= 3.0

    def forward(self, z):
        x = torch.relu(self.fc1(z))
        x = self.fc2(x)
        return x

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