Q-VAE: Quantum Variational Autoencoder#

This tutorial demonstrates how to train and evaluate a Quantum Variational Autoencoder (Q-VAE) model. Q-VAE combines a variational autoencoder with a quantum Boltzmann machine, enabling more powerful generative and representation-learning capabilities.

Objectives#

Understand the architecture and working principles of Q-VAE
Train Q-VAE on the MNIST dataset
Perform image reconstruction and generation
Use Q-VAE for representation learning and classification
Visualize the latent space with t-SNE

Runtime Environment#

Example location: example/qvae_mnist/

train_qvae.ipynb: train the Q-VAE model
train_qvae_classifier.ipynb: representation learning & classification

Dependencies:

pip install torchvision==0.22.0 torchmetrics[image]

1. QVAE Principles at a Glance#

QVAE (Quantum Variational Autoencoder) is a generative model that brings quantum generative models into the latent space of a Variational Autoencoder (VAE). Its core idea is:

> Replace the prior distribution in a conventional VAE with a Quantum Boltzmann Machine (QBM), thereby building a latent-variable model endowed with quantum generative power.

Model Structure#

QVAE comprises the following key components:

Encoder maps input data $\mathbf{x}$ to an approximate posterior distribution $q_\phi(\mathbf{z}|\mathbf{x})$ over latent variables, typically parameterized by a neural network.
Prior models the prior distribution over latent variables $\mathbf{z}$ with a Quantum Boltzmann Machine (QBM). Its Hamiltonian is:

$\mathcal{H}_\theta = \sum_l \Gamma_l \sigma_l^x + \sum_l h_l \sigma_l^z + \sum_{l<m} W_{lm} \sigma_l^z \sigma_m^z$
Decoder maps latent variables $\mathbf{z}$ (or their continuous relaxation $\boldsymbol{\zeta}$ ) back to the data space and reconstructs the original data via:

$p_\theta(\mathbf{x} | \boldsymbol{\zeta}) \sim \text{Bernoulli}(f_\theta(\boldsymbol{\zeta}))$

Training Objective: Q-ELBO#

QVAE maximizes an approximate log-likelihood via a Quantum Evidence Lower BOund (Q-ELBO):

$\mathcal{L}_{\text{Q-ELBO}} = \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})} [\log p_\theta(\mathbf{x} | \boldsymbol{\zeta})] - \tilde{H}(q_\phi(\mathbf{z}|\mathbf{x}) \| p_\theta(\mathbf{z}))$

QBM Sampling & Training#

Positive phase: sample $\mathbf{z} \sim q_\phi(\mathbf{z}|\mathbf{x})$ from the encoder
Negative phase: sample $\mathbf{z} \sim p_\theta(\mathbf{z})$ from the QBM using Monte-Carlo methods or a quantum annealer

Energy serves as the objective; gradients are computed from positive- and negative-phase samples.

2. Model Architecture#

The Encoder and Decoder modules for the auto-encoder architecture are defined, both inheriting from nn.Module.

They are structurally symmetric: each contains a fully-connected layer, LayerNorm, and a tanh activation, and supports L2 weight-decay regularization. The encoder maps high-dimensional inputs to a low-dimensional latent space, while the decoder attempts to reconstruct the original input from the latent representation. Each module exposes a get_weight_decay method to explicitly add weight regularization to the training loss, improving generalization and mitigating over-fitting.

2.1 Encoder#

class Encoder(nn.Module):

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        latent_dim: int,
        weight_decay: float = 0.01,
    ) -> None:
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim
        self.weight_decay = weight_decay
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.norm1 = nn.LayerNorm(hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, latent_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = self.norm1(x)
        x = F.tanh(x)
        x = self.fc2(x)
        return x

    def get_weight_decay(self) -> torch.Tensor:
        """计算权重的L2正则化损失

        对权重矩阵施加L2正则化可以提高模型的泛化能力。

        Returns:
            torch.Tensor: L2正则化损失值
        """
        return self.weight_decay * (
            torch.sum(self.fc1.weight**2) + torch.sum(self.fc2.weight**2)
        )

2.2 Decoder#

class Decoder(nn.Module):

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        latent_dim: int,
        weight_decay: float = 0.01,
    ) -> None:
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim
        self.weight_decay = weight_decay
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.norm1 = nn.LayerNorm(hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, latent_dim)

    def forward(self, z):
        z = self.fc1(z)
        z = self.norm1(z)
        z = F.tanh(z)
        z = self.fc2(z)

        return z

    def get_weight_decay(self) -> torch.Tensor:
        """计算权重的L2正则化损失

        对权重矩阵施加L2正则化可以提高模型的泛化能力。

        Returns:
            torch.Tensor: L2正则化损失值
        """
        return self.weight_decay * (
            torch.sum(self.fc1.weight**2) + torch.sum(self.fc2.weight**2)
        )

2.3 Complete Q-VAE Model#

Refer to the QVAE class in the module manual.

3. Data Preparation#

This function encapsulates loading and pre-processing of the MNIST dataset, returning training and test DataLoader objects.

Data are converted to tensors via ToTensor and 28×28 images are flattened into 784-D vectors by a custom flatten_tensor to suit fully-connected inputs. The training loader shuffles data, while the test loader keeps order for consistent evaluation.

def setup_data_loaders(root, download=True, batch_size=256, use_cuda=False):
   """
   设置MNIST数据集的数据加载器

   Args:
      root (str): 数据存储根目录
      download (bool): 如果数据不存在是否下载，默认为True
      batch_size (int): 每个批次的样本数量，默认为128
      use_cuda (bool): 是否使用GPU，决定是否启用pin_memory优化

   Returns:
      tuple: (train_loader, test_loader) 训练和测试数据加载器
   """
   # 数据预处理
   transform = transforms.Compose([
      transforms.ToTensor(),             # 转换为Tensor
      transforms.Lambda(flatten_tensor)  # 展平：将28x28图像展平成784维向量
      # 等效于：x.reshape(-1) 或 x.flatten()
   ])

   # 加载训练集
   train_set = datasets.MNIST(
      root=root,           # 数据存储路径
      train=True,          # 加载训练集（共60000个样本）
      transform=transform, # 应用定义的数据变换
      download=download    # 如果数据不存在则自动下载
   )

   # 加载测试集
   test_set = datasets.MNIST(
      root=root,           # 数据存储路径
      train=False,         # 加载测试集（共10000个样本）
      transform=transform  # 应用相同的数据变换
   )

   # 数据加载器配置参数
   # 根据是否使用GPU选择不同的优化参数
   # 将num_workers设为0避免多进程问题
   kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {'num_workers': 0}

   # 创建训练数据加载器
   train_loader = DataLoader(
      dataset=train_set,     # 训练数据集
      batch_size=batch_size, # 每个批次的样本数
      shuffle=True,          # 每个epoch打乱数据顺序，防止模型记忆顺序
      **kwargs               # 解包上述配置参数
   )

   # 创建测试数据加载器
   test_loader = DataLoader(
      dataset=test_set,      # 测试数据集
      batch_size=batch_size, # 批次大小（通常与训练集相同）
      shuffle=False,         # 测试集不需要打乱，保证可重复性
      **kwargs               # 解包配置参数
   )

   return train_loader, test_loader

4. Model Training#

This function implements the complete training pipeline of a Quantum Variational Autoencoder (Q-VAE) on MNIST. The model combines classical neural encoder/decoder networks with a Restricted Boltzmann Machine (RBM), optimized by minimizing the negative ELBO loss with weight decay and a KL divergence term that aligns the latent distribution with the prior. All loss components are logged and periodically saved to disk.

def train_qvae(
    train_loader,  # 用于训练QVAE
    device,
    input_dim=784,  # 图片拉伸后的维度
    hidden_dim=512,  # fc1压缩后的维度
    latent_dim=256,  # 隐变量维度， num_visible + num_hidden
    num_var1=128,  # RBM可见层维度
    num_var2=128,  # RBM藏层维度
    dist_beta=10,  # 重叠分布的beta
    weight_decay=0.01,
    batch_size=256,
    epochs=20,
    lr=1e-3,
    kl_beta=0.000001,
    save_path="./models/",
    sampler_type="sa",
):
    # 创建模型
    model, optimizer = create_model(
        train_loader,
        input_dim=input_dim,
        hidden_dim=hidden_dim,
        latent_dim=latent_dim,
        weight_decay=weight_decay,
        dist_beta=dist_beta,
        num_var1=num_var1,
        num_var2=num_var2,
        lr=lr,
        device=device,
        sampler_type=sampler_type,
    )

    # 训练循环
    loss_history = []
    elbo_history = []
    kl_history = []
    cost_history = []

    model.train()  # 设置模型为训练模式
    for epoch in tqdm(range(1, epochs + 1), desc="Training QVAE"):  # 遍历每个训练轮次
        # 训练一个epoch
        avg_loss, avg_elbo, avg_kl, avg_cost = _train_epoch(
            model, train_loader, optimizer, kl_beta, device
        )

        # 记录历史指标
        loss_history.append(avg_loss)
        elbo_history.append(avg_elbo)
        kl_history.append(avg_kl)
        cost_history.append(avg_cost)

        save_list_to_txt(os.path.join(save_path, "loss_history.txt"), loss_history)
        save_list_to_txt(os.path.join(save_path, "elbo_history.txt"), elbo_history)
        save_list_to_txt(os.path.join(save_path, "cost_history.txt"), cost_history)
        save_list_to_txt(os.path.join(save_path, "kl_history.txt"), kl_history)

        # # 保存当前轮次的模型参数
        # model_save_path = os.path.join(save_path, f'davepp_epoch{epoch}.pth')
        # torch.save(model.state_dict(), model_save_path)

        # 打印本轮训练结果
        print(
            f"Epoch {epoch}/{epochs}: "
            f"Loss: {avg_loss:.4f}, "
            f"elbo: {avg_elbo:.4f}, "
            f"KL: {avg_kl:.4f}, "
            f"Cost: {avg_cost:.4f}"
        )

    # 保存模型
    model_save_path = os.path.join(save_path, f"qvae_mnist.pth")
    torch.save(model.state_dict(), model_save_path)
    return model

5. Visualization & Evaluation#

This section supplies two key visualization tools: (1) plot_training_curves for monitoring convergence by plotting training/validation loss and accuracy curves; (2) t_SNE for dimensionality reduction and visualization of latent representations extracted by the QVAE, revealing how different classes are distributed in latent space. Both support automatic saving of high-resolution images and optional real-time display for convenient experiment tracking and reporting.

5.1 Training-Process Visualization#

def plot_training_curves(
    train_loss_history,
    val_loss_history,
    train_acc_history,
    val_acc_history,
    save_path=None,
    show=True,
):
    """
    绘制训练和验证的损失及准确率曲线

    Args:
        train_loss_history: 训练损失历史
        val_loss_history: 验证损失历史
        train_acc_history: 训练准确率历史
        val_acc_history: 验证准确率历史
        save_path: 图像保存路径
    """
    plt.figure(figsize=(12, 5))

    # 损失曲线
    plt.subplot(1, 2, 1)
    plt.plot(train_loss_history, label="Training Loss", color="blue", alpha=0.7)
    plt.plot(val_loss_history, label="Validation Loss", color="red", alpha=0.7)
    plt.title("Training and Validation Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()
    plt.grid(True, alpha=0.3)

    # 准确率曲线
    plt.subplot(1, 2, 2)
    plt.plot(train_acc_history, label="Training Accuracy", color="blue", alpha=0.7)
    plt.plot(val_acc_history, label="Validation Accuracy", color="red", alpha=0.7)
    plt.title("Training and Validation Accuracy")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy (%)")
    plt.legend()
    plt.grid(True, alpha=0.3)

    plt.tight_layout()
    # plt.close()

    # 自动保存
    if save_path is None:
        # 生成默认保存路径
        # timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        save_path = f"results/mlp_training_curves_{timestamp}.png"

    plt.savefig(save_path, dpi=300, bbox_inches="tight")
    print(f"Training curves saved to: {save_path}")
    plt.show()

    if show:
        plt.show()
    else:
        plt.close()  # 不显示时关闭图像，节省内存

5.3 Latent-Space Visualization#

def t_SNE(
    test_loader,
    qvae_model,
    use_std=True,
    point_size=20,
    alpha=0.6,
    epochs=None,
    save_path=None,
    show=True,
):
    """
    QVAE版本的t-SNE可视化

    Args:
        test_loader: 测试数据加载器
        qvae_model: QVAE模型
        use_std: 是否使用标准差
        point_size: 点大小
        alpha: 透明度
        epochs: 训练轮数
        save_path: 保存路径
        show: 是否显示图像
    """
    features = []
    labels = []

    qvae_model.eval()
    device = next(qvae_model.parameters()).device

    with torch.no_grad():
        for batch_idx, (example_data, example_targets) in enumerate(test_loader):
            example_data = example_data.to(device)

            # QVAE前向传播 - 获取潜变量zeta
            _, _, _, zeta = qvae_model(example_data)

            zeta_np = zeta.cpu().numpy()

            for idx in range(zeta_np.shape[0]):
                features.append(zeta_np[idx])
                labels.append(example_targets[idx].item())

    # 创建DataFrame
    feat_cols = [f"dim_{i}" for i in range(zeta_np.shape[1])]
    df = pd.DataFrame(features, columns=feat_cols)
    df["label"] = labels
    df["label"] = df["label"].apply(lambda i: str(i))

    print(f"Extracted {len(features)} samples with {zeta_np.shape[1]} dimensions")

    # 执行t-SNE
    print("Running t-SNE...")
    tsne = TSNE(n_components=2, verbose=1, perplexity=30, max_iter=300, random_state=42)
    tsne_results = tsne.fit_transform(df[feat_cols].values)

    df_tsne = df.copy()
    df_tsne["x-tsne"] = tsne_results[:, 0]
    df_tsne["y-tsne"] = tsne_results[:, 1]

    # 可视化
    # plt.figure(figsize=(10, 8))
    # scatter = plt.scatter(df_tsne['x-tsne'], df_tsne['y-tsne'],
    #                      c=df_tsne['label'].astype(int),
    #                      cmap='tab10', s=point_size, alpha=alpha)
    fig, ax = plt.subplots(figsize=(10, 8))
    scatter = ax.scatter(
        df_tsne["x-tsne"],
        df_tsne["y-tsne"],
        c=df_tsne["label"].astype(int),
        cmap="tab10",
        s=point_size,
        alpha=alpha,
    )

    # 添加颜色条
    cbar = plt.colorbar(scatter, ax=ax, label="Digit")

    # 动态标题和文件名
    training_status = "fully_trained" if epochs and epochs >= 20 else f"epochs_{epochs}"
    title = f"t-SNE Visualization of QVAE Latent Space ({training_status})"
    plt.title(title)
    plt.xlabel("t-SNE dimension 1")
    plt.ylabel("t-SNE dimension 2")
    plt.tight_layout()

    # 自动保存
    if save_path is None:
        # 生成默认保存路径
        # timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        save_path = f"results/t-SNE_QVAE_{training_status}_{timestamp}.png"

    plt.savefig(save_path, dpi=300, bbox_inches="tight")
    print(f"t-SNE plot saved to: {save_path}")
    plt.show()

    if show:
        plt.show()
    else:
        plt.close()  # 不显示时关闭图像，节省内存

    return df_tsne, save_path, training_status

6. Representation Learning & Classification#

Representations learned by Q-VAE can be used for downstream classification:

The function train_mlp_classifier trains a multi-layer perceptron (MLP) classifier whose input features are the representations extracted by the QVAE. It first splits the dataset into training and validation sets, then initializes an MLP model, optimizer, and loss function. During each epoch parameters are updated on the training set and performance is evaluated on the validation set.

def train_mlp_classifier(
    features,
    labels,
    device,
    epochs=100,
    lr=1e-3,
    weight_decay=1e-4,
    batch_size=64,
    seed=42,
    smoke_test=False,
    show=True,
    save_path="./models/",
):
    # 数据分割
    X_train, X_val, y_train, y_val = train_test_split(
        features.numpy(), labels.numpy(), test_size=0.4, random_state=seed
    )

    X_train = torch.FloatTensor(X_train)
    y_train = torch.LongTensor(y_train)
    X_val = torch.FloatTensor(X_val)
    y_val = torch.LongTensor(y_val)

    train_dataset = TensorDataset(X_train, y_train)
    val_dataset = TensorDataset(X_val, y_val)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

    # MLP模型
    mlp = MLP(input_dim=features.shape[1], output_dim=10).to(device)
    optimizer = torch.optim.Adam(mlp.parameters(), lr=lr, weight_decay=weight_decay)
    criterion = nn.CrossEntropyLoss()

    # 记录训练历史
    train_loss_history = []
    val_loss_history = []
    train_acc_history = []
    val_acc_history = []

    # 训练循环
    best_acc = 0
    for epoch in tqdm(range(1, epochs + 1), desc="Training MLP"):
        # 训练阶段
        train_acc, avg_train_loss = _train_mlp_epoch(
            model=mlp,
            data_loader=train_loader,
            optimizer=optimizer,
            criterion=criterion,
            device=device,
        )

        # 验证阶段
        val_acc, avg_val_loss = _eval_mlp_epoch(
            model=mlp, data_loader=val_loader, criterion=criterion, device=device
        )

        # 记录历史
        train_loss_history.append(avg_train_loss)
        val_loss_history.append(avg_val_loss)
        train_acc_history.append(train_acc)
        val_acc_history.append(val_acc)

        # 打印训练和验证指标
        if epoch % 5 == 0:
            print(
                f"Epoch {epoch:3d}: "
                f"Train Loss: {avg_train_loss:.4f}, "
                f"Val Loss: {avg_val_loss:.4f}, "  # 新增验证损失
                f"Train Acc: {train_acc:.2f}%, "
                f"Val Acc: {val_acc:.2f}%"
            )

        if val_acc > best_acc:
            best_acc = val_acc
            model_save_path = os.path.join(save_path, "best_mlp_classifier.pth")
            torch.save(mlp.state_dict(), model_save_path)

    print(f"Best Validation Accuracy: {best_acc:.2f}%")

    # 绘制训练曲线
    curves_save_path = ""
    if not smoke_test:
        curves_save_path = os.path.join(
            save_path, f"mlp_training_curves_epochs_{epochs}.png"
        )
        plot_training_curves(
            train_loss_history=train_loss_history,
            val_loss_history=val_loss_history,
            train_acc_history=train_acc_history,
            val_acc_history=val_acc_history,
            save_path=curves_save_path,
            show=show,
        )
    return mlp, best_acc, curves_save_path

7. Research Applications: QBM-VAE#

The advanced Q-VAE variant QBM-VAE has demonstrated significant value in research:

Single-cell transcriptomics analysis:

Substantially improves clustering accuracy
Detects novel cell sub-types invisible to conventional methods
Provides new leads for target discovery

相关论文：Quantum-Boosted High-Fidelity Deep Learning