概率电路训练

概述

概率电路的训练涉及两个核心问题：参数学习（Parameter Learning）和结构学习（Structure Learning）。与神经网络不同，概率电路的训练需要考虑可追踪推断约束，使得学习问题更具挑战性。¹

本章介绍概率电路的主要训练方法，从经典的最大似然估计到现代的端到端深度学习方法。

---

参数学习

问题形式化

给定数据集 $\mathcal{D}=\{x^{(1)},\dots ,x^{(N)}\}$，其中 $x\in {\mathbb{R}}^d$，概率电路的参数 $\theta$ 包括：

1. 叶节点参数：各输入分布的参数（如高斯分布的均值和方差）
2. 求和节点权重：归一化的混合权重

目标函数：

$$\underset{\theta }{\max }\sum _{n=1}^N\log p_{\theta }(x^{(n)})$$

EM算法

期望最大化（EM）算法是概率电路参数学习的经典方法。

E步骤：计算后验

对于每个求和节点 $S$，计算子节点的后验概率：

$${\gamma }_{S\to i}^{(n)}=P(i\mid x^{(n)},\theta )=\frac{w_i\cdot p_i(x^{(n)})}{{\sum }_j{w}_j\cdot p_j(x^{(n)})}$$

其中 ${\gamma }_{S\to i}^{(n)}$ 表示样本 $n$ 在求和节点 $S$ 处经过子节点 $i$ 的概率。

M步骤：更新权重

更新求和节点的权重：

$$w_i^{new}=\frac{{\sum }_{n=1}^N{\gamma }_{S\to i}^{(n)}}{{\sum }_j{\sum }_{n=1}^N{\gamma }_{S\to j}^{(n)}}$$

M步骤：更新叶分布参数

对于叶节点分布（如高斯分布），更新参数：

$${\mu }_i^{new}=\frac{{\sum }_{n=1}^N{\gamma }_i^{(n)}\cdot x^{(n)}}{{\sum }_{n=1}^N{\gamma }_i^{(n)}}$$ $${\sigma }_i^{2,new}=\frac{{\sum }_{n=1}^N{\gamma }_i^{(n)}\cdot (x^{(n)}-{\mu }_i{)}^2}{{\sum }_{n=1}^N{\gamma }_i^{(n)}}$$

EM算法的收敛性

EM算法保证似然函数单调递增，但可能收敛到局部最优解。

收敛条件：

$$|{\mathcal{L}}^{(t+1)}-{\mathcal{L}}^{(t)}|<ϵ$$

其中 $\mathcal{L}$ 是对数似然。

梯度下降法

现代概率电路训练更常使用梯度下降法，这使其能与深度学习框架（如PyTorch）无缝集成。

损失函数

使用负对数似然作为损失函数：

$$\mathcal{L}(\theta )=-\frac{1}{N}\sum _{n=1}^N\log p_{\theta }(x^{(n)})$$

前向传播

电路的前向传播计算对数似然：

def forward(self, x):
    # 自底向上计算
    log_prob = self.root.evaluate(x)
    return log_prob

反向传播

损失函数对参数的梯度：

$$\frac{\partial \mathcal{L}}{\partial w_i}=-\frac{1}{N}\sum _{n=1}^N\frac{\partial \log p_{\theta }(x^{(n)})}{\partial w_i}$$

对于求和节点 $S$，梯度为：

$$\frac{\partial p_S}{\partial w_i}=p_i(x)$$

梯度计算的特殊性

概率电路的梯度计算有独特之处：

1. 硬 vs 软注意力

求和节点的梯度可以类比为软注意力机制：

$$\frac{\partial \log p_S}{\partial w_i}=\frac{p_i(x)}{{\sum }_j{w}_j{p}_j(x)}-\frac{p_i(x)}{w_i}$$

第一项是正向传播的注意力权重，第二项是权重归一化的修正项。

2. 梯度裁剪

为防止数值不稳定，需要对梯度进行裁剪：

def compute_gradient(self, x):
    grad = self.backward(x)
    # 梯度裁剪
    grad = torch.clamp(grad, min=-10.0, max=10.0)
    return grad

3. 温度缩放

使用温度参数控制分布的平滑度：

$$p_i^{(T)}(x)\propto \exp \left(\frac{\log p_i(x)}{T}\right)$$

• $T>1$：分布更平滑，有利于早期训练
• $T\to 0$：趋向硬最大，适合后期微调

---

结构学习

问题形式化

结构学习旨在同时学习电路的拓扑结构和参数。

目标：

$$\underset{G,\theta }{\max }\sum _{n=1}^N\log p_{(G,\theta )}(x^{(n)})-\lambda \Vert G\Vert$$

其中 $\Vert G\Vert$ 是结构复杂度惩罚项。

自顶向下的贪婪增长

经典的结构学习方法从单根节点开始，逐步添加新节点。

算法流程

1. 初始化：创建空的根节点
2. 循环直到收敛：
   a. 评估当前结构的似然
   b. 尝试多种扩展操作：
      - 添加 Sum 节点（增加混合物数量）
      - 添加 Product 节点（增加独立分解）
      - 扩展现有叶子（增加分布复杂度）
   c. 选择提升最大的操作
   d. 更新结构
3. 评估：使用验证集选择最优结构

扩展操作详解

添加 Sum 节点：

    [Sum]                  [Sum]
      |        →          /     \
   [Child]              [Sum]  [Child]
                       /     \
                   [Ch₁]   [Ch₂]

添加 Product 节点：

    [Sum]                  [Sum]
      |        →          /     \
   [Child]             [Prod]
                     /      \
                 [Child₁] [Child₂]

基于聚类的结构学习

利用数据聚类来指导结构学习是常见方法：

Chow-Liu树扩展

基础版本：将电路根节点建模为变量的树结构。

1. 计算每对变量的互信息 $I(X_i;X_j)$
2. 构建最大生成树
3. 将树编译为电路

RAT-SPN

Randomized Adaptive Tree SPN (RAT-SPN) 方法：

def learn_rat_spn(data, depth, num_subspaces):
    if depth == 0:
        # 叶子节点：单一变量分布
        return LeafNode(data)
    
    # 分组变量
    scopes = partition_variables(data.shape[1], num_subspaces)
    
    # 递归构建
    sum_node = SumNode()
    for scope in scopes:
        child = learn_rat_spn(data[:, scope], depth - 1, num_subspaces)
        sum_node.add_child(child, weight=1.0/len(scopes))
    
    return sum_node

从数据并行学习

PSDD（Probabilistic Sentential Decision Diagram）

PSDD是一种特殊的概率电路，专门用于离散数据的结构学习：

1. vtree构建：建立变量的层次划分
2. 电路学习：在vtree约束下学习电路

LearnPSDD算法

def learn_psdd(vtree, data):
    # 自底向上学习
    for node in vtree.postorder():
        if node.is_leaf():
            return learn_terminal(node, data)
        else:
            # 学习逻辑AND节点
            left = learn_psdd(node.left, data)
            right = learn_psdd(node.right, data)
            return learn_and_node(node, left, right)

端到端结构学习

现代方法将结构学习和参数学习统一为端到端的优化问题。

可微结构搜索

使用Gumbel-Softmax等技巧使结构选择可微：

class DifferentiableCircuit(nn.Module):
    def __init__(self, num_layers):
        super().__init__()
        # 可学习的路由权重
        self.route_weights = nn.Parameter(torch.randn(num_layers, 2))
    
    def forward(self, x, temperature=1.0):
        # Gumbel-Softmax 路由
        logits = self.route_weights / temperature
        probs = F.gumbel_softmax(logits, tau=temperature, hard=False)
        
        # 根据路由选择路径
        return self.routed_forward(x, probs)

---

与深度学习的结合

Einsum Networks (EiNet)

EiNet使用爱因斯坦求和约定高效实现概率电路：

class EinsumNetwork(nn.Module):
    def __init__(self, num_distributions, num_sums):
        super().__init__()
        # 参数化分布
        self.dist_params = nn.Parameter(torch.randn(num_distributions, d))
        
        # 求和节点权重
        self.sum_weights = nn.Parameter(torch.randn(num_sums))
        
        # 使用einops进行高效计算
        self.op = EinsumOp(equation)
    
    def forward(self, x):
        # 计算叶子分布
        leaf_values = self.compute_leaves(x)  # [..., num_distributions]
        
        # Einsum 操作
        output = self.op(leaf_values, self.sum_weights)
        return output

优势：

• GPU加速的批量计算
• 内存高效的稀疏计算
• 与PyTorch无缝集成

神经概率电路

神经概率电路（Neural Probabilistic Circuits，NPC）将神经网络的思想引入概率电路：

可微分的电路结构

class NeuralProbabilisticCircuit(nn.Module):
    def __init__(self, circuit_template):
        super().__init__()
        self.template = circuit_template
        
        # 神经网络参数化电路组件
        self.weight_net = nn.Sequential(
            nn.Linear(d, 64),
            nn.ReLU(),
            nn.Linear(64, num_weights)
        )
        
        self.dist_net = nn.Sequential(
            nn.Linear(d, 64),
            nn.ReLU(),
            nn.Linear(64, num_params_per_dist)
        )
    
    def forward(self, x):
        # 数据依赖的参数生成
        weights = F.softmax(self.weight_net(x), dim=-1)
        dist_params = self.dist_net(x)
        
        # 在参数化电路上评估
        return self.template.evaluate(x, weights, dist_params)

优势

1. 数据依赖的参数：电路参数可以根据输入动态调整
2. 端到端学习：整个系统可微，可以端到端训练
3. 结合两种范式：概率电路的可解释性 + 神经网络的表达能力

---

正则化与泛化

正则化方法

1. 权重衰减

$${\mathcal{L}}_{reg}=\mathcal{L}+\lambda \sum _i{w}_i^2$$

2. Dropout

在求和节点处应用Dropout：

def sum_node_forward(children_values, weights, dropout_rate=0.0):
    output = sum(w * v for w, v in zip(weights, children_values))
    
    if training and dropout_rate > 0:
        mask = torch.rand_like(output) > dropout_rate
        output = output * mask / (1 - dropout_rate)
    
    return output

3. KL散度惩罚

鼓励权重接近均匀分布：

$${\mathcal{L}}_{KL}=\mathcal{L}+\beta \cdot \sum _i{w}_i\log (w_i/K)$$

泛化理论

概率电路的泛化能力与以下因素相关：

因素	影响
电路深度	过深可能导致过拟合
求和节点数量	控制混合物复杂度
结构正则化	偏好简单结构
数据量	充足数据支持复杂结构

---

实践指南

超参数选择

超参数	推荐值	调整策略
学习率	1e-3 ~ 1e-2	初始用较大学习率，逐步衰减
Batch Size	32 ~ 256	根据内存调整
深度	3 ~ 8	数据越复杂，深度可适当增加
Sum节点分支数	2 ~ 16	权衡表达力与效率
Dropout	0.0 ~ 0.3	根据过拟合程度调整

训练稳定性技巧

class StablePCTrainer:
    def __init__(self, circuit, lr=1e-3):
        self.circuit = circuit
        self.optimizer = torch.optim.Adam(circuit.parameters(), lr=lr)
        
        # 学习率调度
        self.scheduler = WarmupCosineScheduler(lr, warmup_steps=100)
    
    def train_step(self, batch):
        # 1. 计算损失
        log_prob = self.circuit(batch)
        loss = -log_prob.mean()
        
        # 2. 反向传播（梯度裁剪）
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.circuit.parameters(), 1.0)
        
        # 3. 更新
        self.optimizer.step()
        self.scheduler.step()
        
        return loss.item()

常见问题与解决

问题	原因	解决方案
数值溢出	概率太小	使用对数空间计算
权重归一化失败	梯度不稳定	添加权重归一化层
过拟合	结构过于复杂	增加正则化、简化结构
训练不收敛	学习率不当	使用学习率调度

---

代码示例：完整训练流程

import torch
import torch.nn as nn
from spn.torch import PCN
 
# 1. 定义电路结构
class SimplePCN(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        # 叶节点：高斯分布参数
        self.loc = nn.Parameter(torch.randn(input_dim, hidden_dim))
        self.scale = nn.Parameter(torch.ones(input_dim, hidden_dim))
        
        # 求和权重
        self.sum_weights = nn.Parameter(torch.ones(hidden_dim) / hidden_dim)
    
    def forward(self, x):
        # 计算高斯分布
        loc = self.loc.unsqueeze(0)  # [1, d, h]
        scale = F.softplus(self.scale).unsqueeze(0)  # [1, d, h]
        
        # 评估
        log_prob = -0.5 * ((x.unsqueeze(-1) - loc) / scale) ** 2
        log_prob = log_prob.sum(dim=1)  # 沿输入维度求和
        
        # 加权和
        log_prob = log_prob + torch.log_softmax(self.sum_weights, dim=0)
        return torch.logsumexp(log_prob, dim=-1)
 
# 2. 训练循环
def train(model, data, epochs=100):
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
    for epoch in range(epochs):
        optimizer.zero_grad()
        log_prob = model(data).mean()
        loss = -log_prob
        loss.backward()
        optimizer.step()
        
        if epoch % 10 == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")
 
# 3. 训练
model = SimplePCN(input_dim=10, hidden_dim=32)
train(model, train_data)

---

总结

概率电路的训练方法涵盖多个层次：

方法	特点	适用场景
EM算法	经典、保证收敛	小规模数据、离线学习
梯度下降	灵活、与DL结合	大规模数据、深度电路
结构学习	自动发现结构	缺乏先验知识
端到端学习	联合优化	现代深度学习场景

现代概率电路训练的关键趋势是与深度学习的深度融合，利用GPU加速、可微分编程等技术支持大规模高效学习。

---

参考

Footnotes

1. Einsum Networks: Fast and Scalable Learning of Tractable Probabilistic Circuits (ICML 2020) ↩