Overview#

Kaiwu-PyTorch-Plugin (KPP)#

Kaiwu-PyTorch-Plugin (KPP) is a PyTorch plugin for training and evaluating Boltzmann Machines (BM) and their restricted variant (RBM) on coherent photonic quantum computers. It uses the Kaiwu SDK to offload Boltzmann-distribution sampling to quantum hardware, while all remaining computations (e.g. parameter updates) stay within the standard PyTorch workflow.

1. Design Goals#

  • Provide native PyTorch programming experience while seamlessly integrating photonic quantum-sampling capabilities.

  • Support definition, training and inference of standard BM/RBM models.

  • Offer extensible interfaces for swapping samplers or energy functions.

2. Features#

Model Support

Standard RBM and fully-connected BM.

Layer Configuration

Visible and hidden layer sizes are customizable.

Sampling Method

Negative-phase sampling is executed on coherent photonic quantum hardware via the Kaiwu SDK.

PyTorch Integration

Model parameters are torch.nn.Parameter objects, enabling automatic differentiation, GPU acceleration and seamless composition with other PyTorch modules.

3. Extension Mechanisms#

  • Energy functions (e.g. Ising form) are decoupled from samplers.

  • Users can replace sampling strategies (e.g. switch to classical MCMC or another backend) by implementing standard interfaces.

  • Custom energy terms are supported for non-standard BM variants.

4. Example Use-Cases#

  • Hand-written digit generation (RBM trained on MNIST)

  • Q-VAE (quantum variational auto-encoder) training pipeline

5. Target Users#

Researchers

Investigate the impact of quantum sampling on energy-based model training.

Developers

Build hybrid classical-quantum generative models.

Educators & Students

Teach Boltzmann-machine principles and hands-on quantum sampling.

6. Typical Workflow#

The typical workflow for training energy-based neural networks with Kaiwu-PyTorch-Plugin is:

  1. Data preparation: load and preprocess training data, convert to model input format

  2. Model definition: instantiate an RBM or BM and set visible/hidden layer sizes

  3. Optimizer setup: use a PyTorch optimizer (SGD, Adam, etc.) to manage model parameters

  4. Training loop:

    • Compute hidden representations from training data (positive phase)

    • Generate samples from the model distribution with the sampler (negative phase)

    • Compute the objective and back-propagate gradients

    • Update model parameters

  5. Model evaluation: use the trained model for feature extraction, classification or generation tasks

 import torch
 from torch.optim import SGD
 from kaiwu.torch_plugin import RestrictedBoltzmannMachine
 from kaiwu.classical import SimulatedAnnealingOptimizer

 # 1. 准备数据
 x = torch.randint(0, 2, (batch_size, num_visible)).float()

 # 2. 定义模型
 rbm = RestrictedBoltzmannMachine(num_visible, num_hidden)

 # 3. 配置优化器和采样器
 optimizer = SGD(rbm.parameters(), lr=0.01)
 sampler = SimulatedAnnealingOptimizer()

 # 4. 训练循环
 for epoch in range(num_epochs):
     h = rbm.get_hidden(x)           # 正相：计算隐藏层
     s = rbm.sample(sampler)         # 负相：模型采样

     optimizer.zero_grad()
     loss = rbm.objective(h, s)      # 计算目标函数
     loss.backward()                 # 反向传播
     optimizer.step()                # 更新参数

7. Citation#

If Kaiwu-PyTorch-Plugin is useful for your research, please cite:

 @software{KaiwuPyTorchPlugin,
     title = {Kaiwu-PyTorch-Plugin},
     author = {{QBoson Inc.}},
     year = {2024},
     url = {https://github.com/QBoson/Kaiwu-pytorch-plugin}
 }

Related papers:

 @article{QuantumBoostedDeepLearning,
     title = {Quantum-Boosted High-Fidelity Deep Learning},
     author = {{QBoson Research Team}},
     year = {2025},
     url = {https://arxiv.org/pdf/2508.11190}
 }
On this page