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 modeltrain_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 with quantum generative capability.
Model Structure#
QVAE comprises the following key components:
Encoder maps input data to an approximate posterior distribution over latent variables, typically parameterized by a neural network.
Prior models the prior distribution over latent variables with a Quantum Boltzmann Machine (QBM). Its Hamiltonian is:
Decoder maps latent variables (or their continuous relaxation ) back to the data space and reconstructs the original data via:
Training Objective: Q-ELBO#
QVAE maximizes an approximate log-likelihood via a Quantum Evidence Lower Bound (Q-ELBO):
QBM Sampling & Training#
Positive phase: sample from the encoder
Negative phase: sample 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 BasicEncoder(Network):
"""Encoder with linear layers and activation function."""
def __init__(self, weight_decay=0.0, **kwargs):
super().__init__(**kwargs)
self.weight_decay = weight_decay
def forward(self, x):
"""Forward pass: encode input."""
return self.encode(x)
def encode(self, x):
"""Encode input through layers."""
logger.debug("encode")
for layer in self._layers:
if self._activation_fct:
x = self._activation_fct(layer(x))
else:
x = layer(x)
return x
def decode(self, x):
"""Decode not implemented for encoder."""
raise NotImplementedError("Decoder not implemented for encoder")
def get_weight_decay(self) -> torch.Tensor:
"""
Compute L2 regularization loss for all linear layers.
Returns:
torch.Tensor: Weight decay loss (0.0 if weight_decay == 0.0).
"""
if self.weight_decay == 0.0:
return torch.tensor(0.0, device=next(self.parameters()).device)
wd = 0.0
for layer in self._layers:
if isinstance(layer, nn.Linear):
wd += torch.sum(layer.weight ** 2)
return self.weight_decay * wd
2.2 Decoder#
class Decoder(BasicDecoder):
"""Alternative decoder using sequential network."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._network = self._create_network()
def _create_network(self):
"""Create sequential network from node_sequence."""
layers = self._node_sequence
module_layers = []
for idx, (n_in, n_out) in enumerate(layers):
module_layers.append(nn.Linear(n_in, n_out))
# Apply activation function: output activation for last layer
act_fct = (self._output_activation_fct
if idx == len(layers) - 1 else self._activation_fct)
module_layers.append(act_fct)
return nn.Sequential(*module_layers)
def decode(self, x):
"""Decode posterior sample."""
logger.debug("Decoder::decode")
return self._network(x)
def encode(self, x):
"""Encode not implemented for decoder."""
raise NotImplementedError("Encoder not implemented for decoder")
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(self, run_tsne=False, compute_energy=False, tsne_interval=10, generate_animation=False):
"""执行完整训练流程"""
logger.info(f"Start training {self.config.type}")
self._setup_data()
self._create_model()
self._setup_tuner()
# 提取测试集标签（用于能量可视化）
y_test = []
for _, labels in self.test_loader:
y_test.extend(labels.numpy().tolist())
y_test = np.array(y_test)
# 调试信息
logger.info(f"Test labels - unique: {np.unique(y_test)}, count: {len(y_test)}")
logger.info(f"Test labels distribution: {np.bincount(y_test)}")
tsne_frames = []
epoch_pbar = tqdm(range(1, self.config.num_epochs + 1), desc="Training Progress")
for epoch in epoch_pbar:
train_loss = self.tuner.train(epoch)
if hasattr(train_loss, 'item'):
train_loss = train_loss.item()
self.train_losses.append(train_loss)
test_loss, input_data, output_data, label_list = self.tuner.test()
if hasattr(test_loss, 'item'):
test_loss = test_loss.item()
self.test_losses.append(test_loss)
epoch_pbar.set_description(f"Epoch {epoch}/{self.config.num_epochs} - Train Loss: {train_loss:.2f}, Test Loss: {test_loss:.2f}")
if epoch % (self.config.num_epochs // 10) == 0 or epoch == self.config.num_epochs:
self._save_reconstruction(epoch, input_data, output_data)
logger.info(f"Epoch {epoch}: Train Loss={train_loss:.4f}, Test Loss={test_loss:.4f}")
if generate_animation and (epoch % tsne_interval == 0 or epoch == self.config.num_epochs):
frame_path = self._save_tsne_frame(epoch)
tsne_frames.append(frame_path)
self.tuner.save_model(config_string=f"final_{self.config.type}")
self._plot_training_curve()
if generate_animation and tsne_frames:
create_tsne_animation(tsne_frames, output_path=self.output_dir)
import shutil
frame_dir = os.path.join(self.output_dir, "temp_tsne_frames")
if os.path.exists(frame_dir):
shutil.rmtree(frame_dir)
if run_tsne:
logger.info("Generating t-SNE visualization...")
self._visualize_tsne()
if compute_energy:
logger.info("Compute BM energy for each sample...")
energies = self.compute_energy(self.model, self.test_loader)
self._save_energy_plot(energies, y_test, output_path=self.output_dir)
logger.info(f"{self.config.type} training completed")
return self.model, self.train_losses, self.test_losses
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.2 潜在空间可视化#
def t_SNE(
test_loader,
qvae_model,
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))
logger.info(f"Extracted {len(features)} samples with {zeta_np.shape[1]} dimensions")
# 执行t-SNE
logger.info("Running t-SNE...")
tsne = TSNE(n_components=2, verbose=1, perplexity=30, max_iter=500, 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]
# 可视化
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}"
training_status = 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")
logger.info(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 fit(self, X, y, validation_split=0.2):
"""训练MLP模型"""
if self.input_dim is None:
self.input_dim = X.shape[1]
logger.info(f"Auto-detected input_dim: {self.input_dim}")
# 数据划分
X_train, X_val, y_train, y_val = train_test_split(
X, y,
test_size=validation_split,
random_state=self.random_state,
stratify=y
)
# 转为 Tensor
X_train = torch.FloatTensor(X_train).to(self.device)
y_train = torch.LongTensor(y_train).to(self.device)
X_val = torch.FloatTensor(X_val).to(self.device)
y_val = torch.LongTensor(y_val).to(self.device)
self.classes_ = np.unique(y)
# 创建数据加载器
train_dataset = TensorDataset(X_train, y_train)
val_dataset = TensorDataset(X_val, y_val)
train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=self.batch_size, shuffle=False)
# 创建模型
self.model = self._create_model()
optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay)
criterion = nn.CrossEntropyLoss()
# 记录训练历史
train_loss_history = []
val_loss_history = []
train_acc_history = []
val_acc_history = []
best_val_acc = 0.0
best_state = None
epoch_pbar = tqdm(range(1, self.epochs + 1), desc="Training MLP")
for epoch in epoch_pbar:
# 训练
train_acc, avg_train_loss = self._train_mlp_epoch(
model=self.model,
data_loader=train_loader,
optimizer=optimizer,
criterion=criterion,
device=self.device
)
# 验证
val_acc, avg_val_loss = self._eval_mlp_epoch(
model=self.model,
data_loader=val_loader,
criterion=criterion,
device=self.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 % 10 == 0:
# print(f"Epoch {epoch}: Train Loss={train_loss:.4f}, Val Loss={val_loss:.4f}, Val Acc={val_acc:.2f}%")
logger.info(
f"Epoch {epoch}/{epoch}: "
f"Train Loss={avg_train_loss:.4f}, Train Acc={train_acc:.2f}% | "
f"Val Loss={avg_val_loss:.4f}, Val Acc={val_acc:.2f}%"
)
# 选择最佳模型
if val_acc > best_val_acc:
best_val_acc = val_acc
best_state = self.model.state_dict()
if self.save_path is not None:
model_save_path = os.path.join(self.save_path, "best_mlp_classifier.pth")
torch.save(best_state, model_save_path)
# 加载最佳模型
self.model.load_state_dict(best_state)
logger.info(f"Best Validation Accuracy: {best_val_acc:.2f}%")
# 绘制训练曲线
if self.save_path is not None:
curves_save_path = os.path.join(
self.save_path, f"mlp_training_curves_epochs_{self.epochs}.png"
)
else:
curves_save_path = None # let plot_training_curves generate default
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=True,
)
return self
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
Related paper: Quantum-Boosted High-Fidelity Deep Learning