变分自编码器（Variational Autoencoder，VAE）是Kingma与Welling于2013年提出的深度生成模型，通过引入概率编码器和重参数化技巧，学习数据的潜在分布并生成新样本。与确定性自编码器不同，VAE在隐空间施加先验分布约束，使采样点都能解码为有效样本。本文用Java实现VAE核心架构，涵盖编码器、解码器、KL散度与重参数化技巧。
核心原理
VAE的核心思想是将传统自编码器的确定性编码替换为概率分布参数化：
- 概率编码器：输入数据x经过神经网络输出潜在变量z的分布参数（均值μ和对数方差log σ²），而非直接输出z。
- 重参数化技巧：从标准正态分布采样ε，通过z = μ + σ·ε得到服从N(μ, σ²)的样本。这一技巧使梯度可以反向传播经过随机采样节点。
- 证据下界（ELBO）：优化目标由重构损失（输入与输出的差异）和KL散度（后验分布与先验分布N(0,I)的差异）组成。
完整Java实现
import java.util.*;
/**
* 变分自编码器（VAE）Java实现
* 包含编码器、重参数化采样、解码器与ELBO损失计算
*/
public class VAE {
// 网络结构参数
private final int inputDim; // 输入维度（如MNIST为784）
private final int hiddenDim; // 隐藏层维度
private final int latentDim; // 潜在空间维度
// 编码器参数：input -> hidden -> (mu, logVar)
private double[][] encW1; // [inputDim][hiddenDim]
private double[] encB1; // [hiddenDim]
private double[][] encW2_mu; // [hiddenDim][latentDim]
private double[] encB2_mu; // [latentDim]
private double[][] encW2_logVar; // [hiddenDim][latentDim]
private double[] encB2_logVar; // [latentDim]
// 解码器参数：latent -> hidden -> input
private double[][] decW1; // [latentDim][hiddenDim]
private double[] decB1; // [hiddenDim]
private double[][] decW2; // [hiddenDim][inputDim]
private double[] decB2; // [inputDim]
private Random random;
public VAE(int inputDim, int hiddenDim, int latentDim) {
this.inputDim = inputDim;
this.hiddenDim = hiddenDim;
this.latentDim = latentDim;
this.random = new Random(42);
initializeParameters();
}
/**
* Xavier初始化所有参数
*/
private void initializeParameters() {
double scale1 = Math.sqrt(2.0 / (inputDim + hiddenDim));
double scale2 = Math.sqrt(2.0 / (hiddenDim + latentDim));
double scale3 = Math.sqrt(2.0 / (latentDim + hiddenDim));
double scale4 = Math.sqrt(2.0 / (hiddenDim + inputDim));
encW1 = randMatrix(inputDim, hiddenDim, scale1);
encB1 = randVector(hiddenDim, scale1);
encW2_mu = randMatrix(hiddenDim, latentDim, scale2);
encB2_mu = randVector(latentDim, scale2);
encW2_logVar = randMatrix(hiddenDim, latentDim, scale2);
encB2_logVar = randVector(latentDim, scale2);
decW1 = randMatrix(latentDim, hiddenDim, scale3);
decB1 = randVector(hiddenDim, scale3);
decW2 = randMatrix(hiddenDim, inputDim, scale4);
decB2 = randVector(inputDim, scale4);
}
private double[][] randMatrix(int rows, int cols, double scale) {
double[][] m = new double[rows][cols];
for (int i = 0; i < rows; i++) {
for (int j = 0; j < cols; j++) {
m[i][j] = random.nextGaussian() * scale;
}
}
return m;
}
private double[] randVector(int len, double scale) {
double[] v = new double[len];
for (int i = 0; i < len; i++) {
v[i] = random.nextGaussian() * scale;
}
return v;
}
/**
* 编码器：输入x，输出潜在分布参数mu和logVar
*/
public EncoderOutput encode(double[] x) {
// 第一层：input -> hidden
double[] h = linear(x, encW1, encB1);
h = relu(h);
// 第二层分支：mu和logVar
double[] mu = linear(h, encW2_mu, encB2_mu);
double[] logVar = linear(h, encW2_logVar, encB2_logVar);
return new EncoderOutput(mu, logVar);
}
/**
* 重参数化技巧：从N(mu, sigma^2)采样z
* z = mu + sigma * epsilon, epsilon ~ N(0, I)
*/
public double[] reparameterize(double[] mu, double[] logVar) {
double[] z = new double[latentDim];
for (int i = 0; i < latentDim; i++) {
double sigma = Math.exp(0.5 * logVar[i]);
double epsilon = random.nextGaussian();
z[i] = mu[i] + sigma * epsilon;
}
return z;
}
/**
* 解码器：从潜在变量z重构输入
*/
public double[] decode(double[] z) {
// 第一层：latent -> hidden
double[] h = linear(z, decW1, decB1);
h = relu(h);
// 第二层：hidden -> input (sigmoid输出概率)
double[] xHat = linear(h, decW2, decB2);
xHat = sigmoid(xHat);
return xHat;
}
/**
* 前向传播：编码 -> 重参数化 -> 解码
*/
public VAEResult forward(double[] x) {
EncoderOutput enc = encode(x);
double[] z = reparameterize(enc.mu, enc.logVar);
double[] xHat = decode(z);
return new VAEResult(xHat, enc.mu, enc.logVar, z);
}
/**
* 计算ELBO损失 = 重构损失 + KL散度
* 对于二值数据（如MNIST），重构损失使用二元交叉熵
*/
public double computeLoss(double[] x, double[] xHat, double[] mu, double[] logVar) {
// 重构损失：二元交叉熵（逐元素）
double reconLoss = 0.0;
for (int i = 0; i < inputDim; i++) {
double p = Math.max(1e-10, Math.min(1 - 1e-10, xHat[i]));
reconLoss -= x[i] * Math.log(p) + (1 - x[i]) * Math.log(1 - p);
}
// KL散度：D_KL(N(mu, sigma^2) || N(0, I))
// = -0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
double klLoss = 0.0;
for (int i = 0; i < latentDim; i++) {
klLoss -= 0.5 * (1 + logVar[i] - mu[i] * mu[i] - Math.exp(logVar[i]));
}
return reconLoss + klLoss;
}
/**
* 从先验分布N(0, I)采样生成新数据
*/
public double[] generate() {
double[] z = new double[latentDim];
for (int i = 0; i < latentDim; i++) {
z[i] = random.nextGaussian();
}
return decode(z);
}
/**
* 在潜在空间插值生成过渡样本
*/
public double[][] interpolate(double[] x1, double[] x2, int steps) {
EncoderOutput enc1 = encode(x1);
EncoderOutput enc2 = encode(x2);
double[] z1 = reparameterize(enc1.mu, enc1.logVar);
double[] z2 = reparameterize(enc2.mu, enc2.logVar);
double[][] results = new double[steps][];
for (int i = 0; i < steps; i++) {
double t = (double) i / (steps - 1);
double[] z = lerp(z1, z2, t);
results[i] = decode(z);
}
return results;
}
// ============ 辅助方法 ============
private double[] linear(double[] x, double[][] w, double[] b) {
int outDim = w[0].length;
double[] out = new double[outDim];
for (int j = 0; j < outDim; j++) {
double sum = b[j];
for (int i = 0; i < x.length; i++) {
sum += x[i] * w[i][j];
}
out[j] = sum;
}
return out;
}
private double[] relu(double[] x) {
double[] out = new double[x.length];
for (int i = 0; i < x.length; i++) {
out[i] = Math.max(0, x[i]);
}
return out;
}
private double[] sigmoid(double[] x) {
double[] out = new double[x.length];
for (int i = 0; i < x.length; i++) {
out[i] = 1.0 / (1.0 + Math.exp(-x[i]));
}
return out;
}
private double[] lerp(double[] a, double[] b, double t) {
double[] out = new double[a.length];
for (int i = 0; i < a.length; i++) {
out[i] = a[i] * (1 - t) + b[i] * t;
}
return out;
}
// ============ 内部类 ============
static class EncoderOutput {
double[] mu;
double[] logVar;
EncoderOutput(double[] mu, double[] logVar) {
this.mu = mu;
this.logVar = logVar;
}
}
static class VAEResult {
double[] xHat; // 重构输出
double[] mu; // 分布均值
double[] logVar; // 分布对数方差
double[] z; // 采样得到的潜在变量
VAEResult(double[] xHat, double[] mu, double[] logVar, double[] z) {
this.xHat = xHat;
this.mu = mu;
this.logVar = logVar;
this.z = z;
}
}
/**
* 训练器：使用SGD优化VAE参数
*/
static class Trainer {
private VAE vae;
private double learningRate;
public Trainer(VAE vae, double learningRate) {
this.vae = vae;
this.learningRate = learningRate;
}
/**
* 简化训练：前向传播并应用数值梯度下降（演示用）
* 实际生产环境应使用自动微分框架
*/
public double trainStep(double[][] batch) {
double totalLoss = 0.0;
for (double[] x : batch) {
VAEResult result = vae.forward(x);
double loss = vae.computeLoss(x, result.xHat, result.mu, result.logVar);
totalLoss += loss;
// 数值梯度更新（简化版，仅更新解码器最后层演示）
// 实际实现应使用反向传播计算所有参数的梯度
double eps = 1e-5;
for (int i = 0; i < vae.decW2.length; i++) {
for (int j = 0; j < vae.decW2[i].length; j++) {
double original = vae.decW2[i][j];
vae.decW2[i][j] = original + eps;
double lossPlus = vae.computeLoss(x, vae.decode(result.z), result.mu, result.logVar);
vae.decW2[i][j] = original - eps;
double lossMinus = vae.computeLoss(x, vae.decode(result.z), result.mu, result.logVar);
vae.decW2[i][j] = original;
double grad = (lossPlus - lossMinus) / (2 * eps);
vae.decW2[i][j] -= learningRate * grad;
}
}
}
return totalLoss / batch.length;
}
}
/**
* 主函数：演示VAE的编码、解码与生成能力
*/
public static void main(String[] args) {
// 模拟MNIST配置：784维输入 -> 256维隐藏 -> 20维潜在空间
int inputDim = 784;
int hiddenDim = 256;
int latentDim = 20;
VAE vae = new VAE(inputDim, hiddenDim, latentDim);
// 构造模拟的二值化MNIST样本（0或1像素）
Random rand = new Random(123);
double[] sample1 = new double[inputDim];
double[] sample2 = new double[inputDim];
for (int i = 0; i < inputDim; i++) {
sample1[i] = rand.nextDouble() > 0.7 ? 1.0 : 0.0;
sample2[i] = rand.nextDouble() > 0.3 ? 1.0 : 0.0;
}
System.out.println("=== VAE编码与解码演示 ===");
// 编码sample1
EncoderOutput enc1 = vae.encode(sample1);
System.out.println("Sample1 编码均值前5维: " +
Arrays.toString(Arrays.copyOf(enc1.mu, 5)));
System.out.println("Sample1 编码logVar前5维: " +
Arrays.toString(Arrays.copyOf(enc1.logVar, 5)));
// 重参数化采样
double[] z1 = vae.reparameterize(enc1.mu, enc1.logVar);
System.out.println("Sample1 潜在向量前5维: " +
Arrays.toString(Arrays.copyOf(z1, 5)));
// 解码重构
double[] recon1 = vae.decode(z1);
double reconError1 = 0.0;
for (int i = 0; i < inputDim; i++) {
reconError1 += Math.abs(sample1[i] - recon1[i]);
}
System.out.printf("Sample1 平均绝对重构误差: %.4f%n", reconError1 / inputDim);
// 计算损失
VAEResult result1 = vae.forward(sample1);
double loss1 = vae.computeLoss(sample1, result1.xHat, result1.mu, result1.logVar);
System.out.printf("Sample1 ELBO Loss: %.4f%n", loss1);
// 从先验生成新样本
System.out.println("\n=== 生成新样本 ===");
double[] generated = vae.generate();
double meanVal = 0.0;
for (double v : generated) meanVal += v;
System.out.printf("生成样本平均像素值: %.4f%n", meanVal / inputDim);
// 潜在空间插值
System.out.println("\n=== 潜在空间插值 ===");
double[][] interpolation = vae.interpolate(sample1, sample2, 5);
for (int i = 0; i < interpolation.length; i++) {
double avg = 0.0;
for (double v : interpolation[i]) avg += v;
System.out.printf("插值步 %d 平均像素值: %.4f%n", i, avg / inputDim);
}
// 多次采样同一输入，展示变分特性
System.out.println("\n=== 变分特性：同一输入多次采样 ===");
for (int trial = 0; trial < 3; trial++) {
double[] z = vae.reparameterize(enc1.mu, enc1.logVar);
double[] recon = vae.decode(z);
double err = 0.0;
for (int i = 0; i < inputDim; i++) err += Math.abs(sample1[i] - recon[i]);
System.out.printf("采样 %d 平均重构误差: %.4f%n", trial + 1, err / inputDim);
}
}
}
关键设计解析
重参数化技巧
这是VAE最核心的创新。直接从N(μ, σ²)采样会导致梯度无法反向传播。通过将随机性外提到标准正态分布ε ~ N(0, I)，采样过程变为确定性变换z = μ + σ·ε，使网络参数μ和σ的梯度可以正常计算。这种”随机节点重参数化”技巧是变分推断与深度网络结合的关键桥梁。
KL散度的解析形式
VAE的优化目标包含KL散度项D_KL(q(z|x) || p(z))。当先验p(z)和后验q(z|x)都是高斯分布时，KL散度有闭式解：
$$D_{KL} = -\frac{1}{2}\sum_{i=1}^{d}(1 + \log(\sigma_i^2) – \mu_i^2 – \sigma_i^2)$$
这使得训练无需蒙特卡洛估计KL项，大幅提升了效率。
潜在空间的结构化
KL散度项充当正则化器，将后验分布拉向标准正态。这确保了潜在空间的连续性：邻近的z值解码为相似的样本。相比之下，传统自编码器的潜在空间可能存在”空洞”——某些区域采样会得到无意义输出。
复杂度分析
| 组件 | 时间复杂度 | 空间复杂度 |
|---|---|---|
| 编码器 | O(inputDim·hiddenDim + hiddenDim·latentDim) | O参数规模 |
| 重参数化 | O(latentDim) | O(latentDim) |
| 解码器 | O(latentDim·hiddenDim + hiddenDim·inputDim) | O参数规模 |
| 前向传播（总体） | O(d·h + h·l) | O(d + h + l) |
其中d为输入维度，h为隐藏维度，l为潜在维度。
扩展方向
- β-VAE：在KL项前添加系数β，调节潜在表示的解耦程度。
- 条件VAE（CVAE）：将类别标签作为额外输入，实现可控生成。
- VAE-GAN：结合VAE的编码器与GAN的判别器，提升生成质量。
- 层次VAE：在潜在空间堆叠多层VAE，学习更复杂的分布结构。
