BM Generation: Boltzmann Machine Data Generation#

This tutorial demonstrates how to perform unsupervised training and data generation with a fully-connected Boltzmann Machine (BM). The method is suitable for quickly producing large numbers of small-scale samples.

Objectives#

Understand the differences between fully-connected and restricted Boltzmann machines
Train a BM using KL divergence and contrastive divergence
Implement learning-rate scheduling and sampling strategies
Visualize the distribution of generated samples

Runtime Environment#

Example location: example/bm_generation/

train_bm.ipynb: training script
sample_bm.ipynb: sampling and testing script

Dependencies:

 pip install kaiwu==1.3.0 pandas matplotlib

1. Introduction to Fully-Connected Boltzmann Machines#

1.1 BM vs RBM#

Feature	Restricted Boltzmann Machine (RBM)	Fully-Connected Boltzmann Machine (BM)
Connectivity	Fully connected between layers, none within layers	All nodes fully connected
Sampling strategy	Parallelizable sampling	Sequential sampling required
Training efficiency	Higher	Lower
Expressive power	Limited by bipartite structure	Stronger, can model arbitrary distributions

1.2 Applicable Scenarios#

Fully-connected Boltzmann machines are suited for:

Scenarios requiring modeling of complex dependencies
Small-scale sample generation
Research on sampling from Boltzmann distributions

1.3 Data-Generation Pipeline#

Negative-phase sampling calls self.bm_net.sample(self.worker) to draw a complete set of visible-hidden states state_all from the current BM joint distribution, approximating the model distribution for contrastive-divergence-like objectives.
Positive-phase sampling performs two kinds of positive sampling for every input batch data:
- Full conditional sampling: fixes the entire input data as visible values and samples the corresponding complete state state_v.
- Partial conditional sampling: fixes only the non-output part of the input (i.e. data[:, :-num_output]) and freely samples the output dimensions, yielding state_vi.
These two samplings are used to compute:
- KL-divergence term (kl_divergence): measures the discrepancy between the model and data distributions.
- Negative conditional-likelihood term (ncl): encourages the model to reconstruct the output part correctly given the input.
Objective construction The final optimization target is the weighted combination:

\mathcal{L} = \alpha \cdot \text{KL\_divergence} + (1 - \alpha) \cdot \text{NCL}

where $\alpha$ is controlled by cost_param["alpha"], balancing generative power and conditional consistency.

Multi-process acceleration splits one batch across workers; each sub-process independently calls process_solve_graph to perform positive-phase sampling and probability estimation, and the results are merged for gradient computation.

1.4 Data-Generation Features#

Conditional generation support: part of the visible units can be fixed as observations while the rest are generated by the model, suitable for semi-supervised or sequence-completion tasks.
Differentiable training: all sampling operations are embedded in the PyTorch computation graph, enabling end-to-end back-propagation.
Visualization support: weight matrices and their gradients can be plotted in real time during training for debugging and analysis.
Flexible output structure: the num_output parameter explicitly separates “input” and “output” visible units, adapting the model to supervised settings such as regression or classification.

2. Loading Data#

CSVDataset is a lightweight dataset wrapper inheriting from PyTorch Dataset, used to load structured data stored as NumPy arrays or lists (e.g. data read from CSV files).

It implements __len__ and __getitem__ so it works seamlessly with DataLoader. Every time a sample is accessed by index, the data are automatically converted to torch.float32 tensors, satisfying neural-network input requirements. The class is suitable for unsupervised/self-supervised tasks or as a generic data-loader base component.

class CSVDataset(Dataset):
    """
    数据集

    Args:
        data: 输入数据，形状为 (numcol, numrow)
    """

    def __init__(self, data):
        self.data = data

    def __len__(self):
        """
        返回数据集的大小
        """
        return len(self.data)

    def __getitem__(self, idx):
        """
        根据索引获取数据
        Args:
            idx: 索引

        Return:
            torch.Tensor: 数据
        """
        return torch.tensor(self.data[idx], dtype=torch.float32)

3. Model Construction#

This initialization method configures a Boltzmann-machine-based training framework. It receives input data, a result saver (saver), and a task handler (worker), and sets the numbers of visible, hidden, and output units. Internally it builds a Boltzmann network (BoltzmannMachine) whose total units equal visible plus hidden units, and initializes learning-related hyper-parameters including momentum and regularization coefficients (alpha and beta).

This structure provides the foundation for subsequent energy modeling, sampling-based training, and model evaluation, and is applicable to unsupervised or generative learning tasks.

    def __init__(
        self, data, saver, worker, num_visible=100, num_hidden=10, num_output=10
    ):
        """初始化 Trainer。

        Args:
            data (list): 训练数据列表，每个元素应为 torch.Tensor。
            saver (object): 具备 save_info 和 output_loss 方法的持久化对象。
            worker (object): 执行采样的后端 worker。
            num_visible (int): 可见层维数。默认为 100。
            num_hidden (int): 隐藏层维数。默认为 10。
            num_output (int): 输出层维数。默认为 10。
        """
        self.data = data
        self.saver = saver
        self.worker = worker
        self.num_visible = num_visible
        self.num_hidden = num_hidden
        self.num_output = num_output

        self.learning_parameters = {
            "learning_rate": 0.001,
            "weight_decay_rate": 0.0,
            "momentum_rate": 0.0,
        }
        self.bm_net = BoltzmannMachine(num_nodes=self.num_visible + self.num_hidden)

        # 核心：将模型参数移动到共享内存中，供多进程直接访问
        if self.bm_net.device.type != "cpu":
            self.bm_net_cpu = BoltzmannMachine(
                num_nodes=self.num_visible + self.num_hidden
            ).to(torch.device("cpu"))
        else:
            self.bm_net_cpu = self.bm_net
        self.bm_net_cpu.share_memory()

        self.cost_param = {"alpha": 0.5, "beta": 0.5}

4. Training Pipeline#

4.1 Learning-Rate Scheduler#

CosineScheduleWithWarmup is a custom learning-rate scheduler inheriting from PyTorch’s LambdaLR. It first applies a linear warmup to gradually increase the learning rate, then uses cosine annealing to smoothly decay it to zero.

This scheduling helps the model converge stably in early stages and fine-tune parameters later, and is widely used in deep-learning tasks to improve training efficacy and generalization. Users can flexibly configure warmup steps, total training steps, and cosine cycles.

class CosineScheduleWithWarmup(LambdaLR):
    """带有warmup的余弦退火学习率调度器"""

    def __init__(
        self,
        optimizer,
        num_warmup_steps,
        num_training_steps,
        num_cycles=0.5,
        last_epoch=-1,
    ):
        self.num_warmup_steps = num_warmup_steps
        self.num_training_steps = num_training_steps
        self.num_cycles = num_cycles
        super(CosineScheduleWithWarmup, self).__init__(
            optimizer, self.lr_lambda, last_epoch
        )

    def lr_lambda(self, current_step):
        """带有warmup的cosine schedule学习率生成器"""
        if current_step < self.num_warmup_steps:
            return float(current_step) / float(max(1, self.num_warmup_steps))
        descent_steps = float(max(1, self.num_training_steps - self.num_warmup_steps))
        progress = float(current_step - self.num_warmup_steps) / descent_steps
        return max(
            0.0,
            0.5 * (1.0 + math.cos(math.pi * float(self.num_cycles) * 2.0 * progress)),
        )

4.2 Trainer Implementation#

This training method implements a custom optimization procedure for Boltzmann machines that combines a hybrid loss of KL divergence and negative conditional likelihood (NCL).

Adam optimizer is used together with the cosine-annealing scheduler with warmup to improve convergence stability. Every step generates global states via sampling, and multi-process parallelism accelerates sub-tasks. Loss values are printed periodically, while model parameters and training information are saved every few steps. Visualization of weights and gradients is also supported for debugging and analyzing training dynamics.

    def train(self, max_steps, save_path, num_processes=1):
        """执行模型训练主循环。

        通过多进程并行采样并结合 Adam 优化器更新玻尔兹曼机参数。

        Args:
            max_steps (int): 最大训练步数。
            save_path (str): 模型和日志的保存路径。
            num_processes (int): 并行采样的进程数。默认为 1。
        """
        optimizer = torch.optim.Adam(
            self.bm_net.parameters(),
            lr=self.learning_parameters["learning_rate"],
            weight_decay=self.learning_parameters["weight_decay_rate"],
        )
        scheduler = CosineScheduleWithWarmup(
            optimizer,
            num_training_steps=max_steps,
            num_warmup_steps=int(max_steps / 20),
            num_cycles=0.5,
        )

        t_start = time.time()
        step = 0
        self.saver.save_info(self.bm_net, save_path, 0, 0.0)

        # 预先分配 Pool 以减少重复创建开销
        pool = mp.Pool(processes=num_processes)

        while step < max_steps:
            for batch_data in self.data:
                if step >= max_steps:
                    break

                optimizer.zero_grad()
                step += 1

                # 1. 负相位采样 (从当前模型分布采样)
                # 这个步骤通常对整个模型进行，直接由主进程执行
                state_all = self.bm_net.sample(self.worker).detach()

                # 2. 多进程并行正相位采样 (共享存储模式)
                # 将本批次数据设为共享内存
                self._sync_gpu_to_cpu()
                cpu_data = batch_data.cpu().share_memory_()

                # 按进程数切分数据
                chunks = torch.chunk(cpu_data, num_processes)

                # 准备子进程参数 (传递共享的模型对象)
                sd_args = [
                    (self.bm_net_cpu, self.worker, chunk, self.num_output)
                    for chunk in chunks
                    if chunk.size(0) > 0
                ]

                # 并行执行采样
                all_results = pool.map(process_solve_graph_worker, sd_args)

                # 3. 汇总结果并计算 Loss
                kl_divergence = torch.tensor(0.0, device=state_all.device)
                ncl = torch.tensor(0.0, device=state_all.device)

                # 收集所有进程的采样状态
                combined_v = torch.cat([res[0] for res in all_results], dim=0).to(
                    self.bm_net.device
                )
                combined_vi = torch.cat([res[1] for res in all_results], dim=0).to(
                    self.bm_net.device
                )

                # 计算 KL 散度项
                kl_divergence = self.bm_net.objective(combined_v, state_all)

                # 计算 NCL (Non-Contrastive Loss) 类似项
                # 保持原逻辑：将输出部分置零后计算 objective
                v_ncl = combined_v.clone()
                vi_ncl = combined_vi.clone()
                v_ncl[:, -self.num_output :] = 0.0
                vi_ncl[:, -self.num_output :] = 0.0
                ncl = self.bm_net.objective(v_ncl, vi_ncl)

                # 组合目标函数
                obj = (
                    self.cost_param["alpha"] * kl_divergence
                    + (1 - self.cost_param["alpha"]) * ncl
                )

                # 4. 反向传播与优化
                obj.backward()

                # 可视化 (保留原逻辑)
                if step % 10 == 0:
                    plt.figure(figsize=(8, 4))
                    plt.subplot(1, 2, 1)
                    plt.imshow(self.bm_net.quadratic_coef.detach().cpu().numpy())
                    plt.title("Weights")
                    plt.subplot(1, 2, 2)
                    plt.imshow(self.bm_net.quadratic_coef.grad.cpu().numpy())
                    plt.title("Gradients")
                    plt.show()

                optimizer.step()
                scheduler.step()

                # 输出与保存
                self.saver.output_loss(
                    step, kl_divergence.item(), ncl.item(), obj.item()
                )

                if step % 10 == 0:
                    t_now = time.time()
                    self.saver.save_info(self.bm_net, save_path, step, t_now - t_start)

        pool.close()
        pool.join()

5. Saving and Loading Models#

The Saver class persists key information during model training.

It provides two functions: (1) periodically saves the current model in PyTorch format with step-numbered files for later loading or resuming; (2) appends the loss of every step to the loss.txt log file, enabling visualization and performance tracking of the training process.

class Saver:
    """用于保存信息"""

    def __init__(self, log_path="./log"):
        """初始化"""
        self.log_path = log_path
        if not os.path.exists(log_path):
            os.makedirs(log_path)

    def save_info(self, model, save_path, output_i, time):
        """保存模型等信息

        Args:
            model (Model): 构造模型相关参数
            bias_path (str): 一次项系数的路径
            interation_path (str): 二次项系数的路径
            t_run (float): 计算运行时间
        """
        save_path = os.path.join(save_path, f"rbm_model{output_i}.pth")
        print(f"time: {time}, save_path: {save_path}")
        torch.save(model, save_path)

    def output_loss(self, output_i, kl_div, ncl, func):
        """输出loss

        Args:
            model (Model): 构造模型相关参数
            pr_vi (float): graph_out_hidden相关概率
            pr_v (float): graph_hidden相关概率
            output_i (int): 计算步数
        """
        print(f"step:{output_i}, kl_div:{kl_div}, ncl:{ncl}, cost:{func}")
        with open(os.path.join(self.log_path, "loss.txt"), "a", encoding="utf8") as f:
            f.write(f"step:{output_i}, kl_div:{kl_div}, ncl:{ncl}, cost:{func}\n")