PyTorch与Numpy带截距的最小二乘法：性能优化问询

背景 #

当前使用Numpy对大规模向量做回归分析，现有数据量需隔夜计算，未来数据量将数倍增长，计划迁移至PyTorch提升性能。需求为求解带截距的最小二乘解：predictions @ x = betas，各变量维度：

• predictions: (750, 6340)
• betas: (750, 4313)
• x: (6340, 4313)（每个列对应predictions中一列的回归系数，不含截距）

当前Numpy实现需循环遍历predictions的每一列，代码如下：

for candidate in range(0, predictions.shape[1]):  # each column is a candidate
    prediction = predictions[:, candidate]
    # allow for an intercept by adding a column with ones
    prediction = np.vstack([prediction, np.ones(prediction.shape[0])]).T
    sol = np.linalg.lstsq(prediction, betas, rcond=-1)

问题1：Numpy中是否存在无需循环的实现方式？ #

有，可通过批量矩阵运算实现，利用广播避免循环，效率远高于逐列处理。核心思路是将所有predictions列的截距扩展操作一次性完成，再用矩阵公式批量计算最小二乘解（由于每个回归仅含1个特征+截距，可直接用2x2矩阵的逆公式，比通用lstsq更快）：

import numpy as np

# 构造带截距的批量设计矩阵
# predictions shape: (750, 6340)
X = np.concatenate([predictions[..., None], np.ones_like(predictions)[..., None]], axis=-1)
# X shape: (750, 6340, 2) → 每个位置对应(特征值, 1)

# 计算X^T X和X^T y
X_T = np.transpose(X, (2, 1, 0))  # shape: (2, 6340, 750)
X_T_X = np.matmul(X_T, X)          # shape: (2, 6340, 2) → 每个候选对应的(2,2)矩阵
X_T_y = np.matmul(X_T, betas[:, None, :])  # shape: (2, 6340, 4313)

# 计算2x2矩阵的逆（比通用lstsq更高效）
a = X_T_X[0, :, 0]
b = X_T_X[0, :, 1]
c = X_T_X[1, :, 0]
d = X_T_X[1, :, 1]
det = a * d - b * c  # 每个候选的矩阵行列式

# 构造逆矩阵
inv_X_T_X = np.stack([
    np.stack([d, -b], axis=-1),
    np.stack([-c, a], axis=-1)
], axis=0) / det[None, :, None]

# 批量计算解
sol = np.matmul(inv_X_T_X, X_T_y)  # shape: (2, 6340, 4313)
x = sol[0, :, :]  # shape: (6340, 4313) → 提取特征系数，即目标矩阵x

若偏好使用np.linalg.lstsq，可将输入整理为批量格式（需注意内存占用）：

# 转换为批量输入格式：(候选数, 样本数, 特征数)
X_batch = np.transpose(X, (1, 0, 2))  # shape: (6340, 750, 2)
betas_batch = np.broadcast_to(betas[None, ...], (6340, 750, 4313))  # 广播避免复制

# 批量求解
sol_batch = np.linalg.lstsq(X_batch, betas_batch, rcond=-1)[0]
x = sol_batch[:, 0, :]  # shape: (6340, 4313)

问题2：使用statsmodels.OLS能否避免循环？ #

难以完全避免循环。statsmodels的OLS接口基于单个设计矩阵，而每个候选对应不同的特征列（即使都是单特征+截距），无法通过单个设计矩阵覆盖所有候选的回归需求。

若坚持使用statsmodels，可通过列表推导式简化循环代码，但本质仍为逐候选处理：

import statsmodels.api as sm

# 预构造截距列
intercept = np.ones(predictions.shape[0])

# 列表推导式批量处理
sol_list = [
    sm.OLS(betas, sm.add_constant(predictions[:, candidate])).fit().params
    for candidate in range(predictions.shape[1])
]

# 拼接得到x矩阵（取每个回归的特征系数，即params[0]）
x = np.vstack([sol[0] for sol in sol_list])  # shape: (6340, 4313)

注意：statsmodels侧重统计推断，计算效率远低于Numpy/PyTorch的批量矩阵运算，不适合大规模数据场景。

问题3：PyTorch中如何添加截距并高效计算？ #

PyTorch支持批量矩阵运算和torch.linalg.lstsq的批量处理，可参考Numpy的思路，优先使用2x2矩阵逆公式提升效率：

import torch

# 转换为PyTorch张量
t_predictions = torch.tensor(predictions, dtype=torch.float)
t_betas = torch.tensor(betas, dtype=torch.float)

# 构造带截距的批量设计矩阵
ones = torch.ones_like(t_predictions)
X = torch.stack([t_predictions, ones], dim=-1)  # shape: (750, 6340, 2)

# 计算X^T X和X^T y
X_T = X.permute(2, 1, 0)  # shape: (2, 6340, 750)
X_T_X = torch.matmul(X_T, X)  # shape: (2, 6340, 2)
X_T_y = torch.matmul(X_T, t_betas.unsqueeze(0))  # shape: (2, 6340, 4313)

# 计算2x2矩阵的逆
a = X_T_X[0, :, 0]
b = X_T_X[0, :, 1]
c = X_T_X[1, :, 0]
d = X_T_X[1, :, 1]
det = a * d - b * c

inv_X_T_X = torch.stack([
    torch.stack([d, -b], dim=-1),
    torch.stack([-c, a], dim=-1)
], dim=0) / det[None, :, None]

# 批量求解
sol = torch.matmul(inv_X_T_X, X_T_y)
x = sol[0, :, :]  # shape: (6340, 4313) → 目标矩阵x

若使用torch.linalg.lstsq批量处理：

# 转换为批量输入格式：(候选数, 样本数, 特征数)
X_batch = X.permute(1, 0, 2)  # shape: (6340, 750, 2)
betas_batch = t_betas.unsqueeze(0).expand(6340, -1, -1)  # 广播

# 批量求解
t_sol = torch.linalg.lstsq(X_batch, betas_batch)
x = t_sol.solution[:, 0, :]  # shape: (6340, 4313)

若使用GPU加速，只需将张量移至GPU（t_predictions = t_predictions.cuda()），其余代码无需修改，可大幅提升计算速度。

内容的提问来源于stack exchange，提问作者Feillen