balance源码解析：核心算法CBPS的Python实现原理

吕奕昶

699人浏览 · 2026-07-12 09:37:17

吕奕昶 · 2026-07-12 09:37:17 发布

balance源码解析：核心算法CBPS的Python实现原理

【免费下载链接】balance The balance python package offers a simple workflow and methods for dealing with biased data samples when looking to infer from them to some target population of interest. 项目地址: https://gitcode.com/gh_mirrors/ba/balance

CBPS（协变量平衡倾向得分）算法是balance包中处理有偏数据样本的核心加权方法。作为Facebook Research开发的开源Python包，balance通过CBPS算法为研究人员提供了强大的偏差校正工具，特别适用于调查统计、观察性研究和因果推断场景。本文将深入解析balance中CBPS算法的实现原理，帮助您理解这一先进统计方法背后的技术细节。

什么是CBPS算法？

CBPS（Covariate Balancing Propensity Score）是一种创新的倾向得分估计方法，由Imai和Ratkovic于2014年提出。与传统倾向得分方法不同，CBPS在估计倾向得分的同时最大化协变量平衡，有效解决了模型误设问题。在balance包中，CBPS算法被实现为处理有偏样本的核心加权方法，位于balance/weighting_methods/cbps.py文件中。

CBPS在偏差校正流程中的位置

CBPS算法的核心数学原理

CBPS基于广义矩方法（GMM）框架，将倾向得分估计问题转化为优化问题。算法的核心思想是同时满足两个条件：

最大似然条件：最大化样本包含的预测准确性
协变量平衡条件：确保加权后样本与目标群体的协变量分布一致

数学上，CBPS通过以下矩条件实现：

$$ g_i(I_i,X_i)=\begin{pmatrix} \frac{I_i-\pi_\beta(X_i)}{\pi_\beta(X_i)(1-\pi_\beta(X_i))}\pi'\beta(X_i) \ \frac{I_i-\pi\beta(X_i)}{\pi_\beta(X_i)(1-\pi_\beta(X_i))} X_i \end{pmatrix} $$

其中$\pi_\beta(X_i)$是倾向得分，$I_i$是样本指示变量，$X_i$是协变量向量。

balance中CBPS的实现架构

1. 核心函数结构

balance的CBPS实现包含多个关键函数模块：

主函数：cbps() - 位于balance/weighting_methods/cbps.py，负责完整的算法流程
GMM损失函数：gmm_function() - 计算广义矩方法的损失和权重矩阵
平衡损失函数：bal_loss() - 专门计算协变量平衡的损失
辅助函数：包括logit_truncated()、compute_pseudo_weights_from_logit_probs()等

2. 数据处理流程

CBPS算法的数据处理流程分为几个关键步骤：

# 1. 数据预处理和变量选择
variables = balance_util.choose_variables(sample_df, target_df, variables=variables)

# 2. 数据转换和模型矩阵构建
model_matrix_output = balance_util.model_matrix(
    sample_df, target_df, variables, add_na=(na_action == "add_indicator")
)

# 3. 矩阵标准化和SVD分解
model_matrix_standardized = _standardize_model_matrix(X_matrix, X_matrix_columns_names)
U, s, Vh = scipy.linalg.svd(X_matrix, full_matrices=False)

3. 优化算法实现

balance采用两阶段GMM估计器，这是对原始论文中连续更新估计器的改进，显著提高了计算速度：

# 第一阶段：GLM初始估计
glm_mod = sm.GLM(in_pop, U, family=sm.families.Binomial(), freq_weights=design_weights)
beta_0 = glm_mod.fit().params

# 第二阶段：GMM优化
rescale_initial_result = scipy.optimize.minimize(
    alpha_function, x0=[1], args=(beta_0, U, design_weights, in_pop)
)

# 第三阶段：平衡损失最小化
balance_optimize_result = scipy.optimize.minimize(
    fun=bal_loss, x0=gmm_init, args=(U, design_weights, in_pop, XtXinv)
)

收入变量在应用CBPS权重前后的QQ图对比

CBPS算法的关键技术创新

1. 矩阵奇异值分解（SVD）处理

balance中的CBPS实现使用了SVD来处理秩不足问题：

# 移除接近零的奇异值
singular_value_threshold = 1e-10
U = U[:, s > singular_value_threshold]
Vh = Vh[s > singular_value_threshold, :]
s = s[s > singular_value_threshold]

# 使SVD符号确定化
U, Vh = sklearn.utils.extmath.svd_flip(U, Vh, u_based_decision=False)

2. 设计效应约束

算法支持通过max_de参数限制权重的设计效应：

if max_de is not None:
    constraints += [
        {
            "type": "ineq",
            "fun": lambda x: (
                max_de - compute_deff_from_beta(U, x, design_weights, in_pop)
            ),
        }
    ]

3. 权重修剪机制

为防止极端权重，balance提供了两种修剪策略：

均值比率修剪：weight_trimming_mean_ratio
百分位数修剪：weight_trimming_percentile

年龄分布在应用CBPS权重前后的对比

CBPS的两种工作模式

1. 过识别模式（"over"）

这是默认模式，同时优化倾向得分和协变量平衡条件：

if cbps_method == "over":
    # 从两个起点进行优化：beta_balance和gmm_init
    gmm_optimize_result_glm_init = scipy.optimize.minimize(
        fun=gmm_loss, x0=gmm_init, args=(U, design_weights, in_pop, invV)
    )
    
    gmm_optimize_result_bal_init = scipy.optimize.minimize(
        fun=gmm_loss, x0=beta_balance, args=(U, design_weights, in_pop, invV)
    )

2. 精确识别模式（"exact"）

仅考虑协变量平衡条件，忽略倾向得分条件：

elif cbps_method == "exact":
    if balance_optimize_result["success"] is np.bool_(False):
        raise Exception("There is no solution satisfying the constraints.")
    beta_opt = beta_balance

实际应用示例

快速开始使用CBPS

from balance.weighting_methods.cbps import cbps
import pandas as pd

# 准备样本和目标数据
sample_df = pd.DataFrame({"age": [25, 30, 35, 40], "income": [50000, 60000, 70000, 80000]})
target_df = pd.DataFrame({"age": [28, 32, 38, 42], "income": [55000, 65000, 75000, 85000]})

# 应用CBPS算法
result = cbps(
    sample_df=sample_df,
    sample_weights=pd.Series([1.0, 1.0, 1.0, 1.0]),
    target_df=target_df,
    target_weights=pd.Series([1.0, 1.0, 1.0, 1.0]),
    variables=["age", "income"],
    cbps_method="over",
    max_de=2.0
)

# 获取调整后的权重
adjusted_weights = result["weight"]

在完整工作流中使用

在balance的完整工作流中，CBPS通常通过BalanceFrame类调用：

from balance import BalanceFrame

# 创建BalanceFrame对象
bf = BalanceFrame(sample_df, sample_weights, target_df, target_weights)

# 使用CBPS进行加权
bf_adjusted = bf.adjust(method="cbps", variables=["age", "income", "gender"])

# 评估平衡效果
bf_adjusted.summary()

性别分布在应用CBPS权重前后的对比

CBPS算法的性能优化

1. 数值稳定性处理

balance的CBPS实现包含多项数值稳定性措施：

倾向得分截断：防止概率值为0或1
奇异值过滤：移除接近零的奇异值
矩阵条件数检查：确保数值计算的稳定性

2. 内存优化

对于大规模数据集，算法使用稀疏矩阵表示：

model_matrix_output = balance_util.model_matrix(
    sample_df, target_df, variables,
    return_type="one",
    return_var_type="sparse",  # 使用稀疏矩阵
    one_hot_encoding=False
)

3. 并行计算支持

虽然当前版本主要依赖单线程优化，但算法设计考虑了未来的并行化扩展。

与其他加权方法的比较

在balance包中，CBPS与其他加权方法共存：

IPW（逆概率加权）：传统的倾向得分方法
Raking：边际平衡方法
Post-stratification：事后分层方法
CBPS：协变量平衡倾向得分（本文重点）

CBPS的主要优势在于：

更好的协变量平衡
对模型误设更稳健
理论性质更优（基于GMM框架）

应用CBPS权重后收入分布与目标群体更加接近

技术实现细节

1. 损失函数计算

gmm_function()函数实现了GMM损失的核心计算：

def gmm_function(beta, X, design_weights, in_pop, invV=None):
    probs = logit_truncated(X, beta)
    N = np.sum(design_weights)
    N_target = np.sum(design_weights[in_pop == 1.0])
    
    weights = compute_pseudo_weights_from_logit_probs(probs, design_weights, in_pop)
    
    # 生成矩条件
    gbar = np.concatenate((
        1.0/N * (np.matmul((X * design_weights[:, None]).T, (in_pop - probs))),
        1.0/N * (np.matmul((X * design_weights[:, None]).T, weights)),
    ))
    
    # 计算损失
    loss = np.matmul(np.matmul(gbar.T, invV), gbar)
    return {"loss": loss, "invV": invV}

2. 权重计算逻辑

最终的权重计算考虑了设计权重和倾向得分：

# 计算最终概率和权重
probs = logit_truncated(U, beta_opt)
weights = np.absolute(
    compute_pseudo_weights_from_logit_probs(probs, design_weights, in_pop)
)

# 应用设计权重
weights = design_weights[in_sample == 1.0] * weights[in_sample == 1.0]

# 权重修剪
weights = balance_adjustment.trim_weights(
    weights, weight_trimming_mean_ratio, weight_trimming_percentile
)

# 归一化到目标大小
weights = weights / original_sum_weights * np.sum(target_weights)