优化Numpy向量运算：多项式逻辑回归梯度下降实现的错误排查与性能提升问询

我正在尝试使用**多项式逻辑回归（multinomial logistic regression）与梯度下降（gradient descent）**训练一个多分类器。具体参数定义：

• 模型权重矩阵 w：形状为 (C, D)，其中C是类别数量，D是输入特征数
• 偏置向量 b：维度为 (C,)
• 训练输入矩阵 X：形状为 (N, D)，N是训练样本数量
• 标签向量 y：形状为 (N,)，每个元素取值范围是0到C-1，对应输入的类别

初始（正确但非向量化）代码 #

这段代码能得到正确的权重与偏置，但未充分利用Numpy运算提升速度：

for _ in range(max_iterations):
    z = np.apply_along_axis(lambda v: v - max(v), 1, X @ w.T + b)
    probs = np.exp(z)
    denom = np.sum(probs, axis=1)
    for i in range(C):
        for j in range(N):
            if i == y[j]:
                w[i] -= (step_size / N) * ((probs[j][i] / denom[j]) - 1) * X[j]
                b[i] -= (step_size / N) * ((probs[j][i] / denom[j]) - 1)
            else:
                w[i] -= (step_size / N) * (probs[j][i] / denom[j]) * X[j]
                b[i] -= (step_size / N) * (probs[j][i] / denom[j])

尝试优化但结果错误的代码 #

这段代码运行速度有所提升，但结果不正确：

for _ in range(max_iterations):
    z = np.apply_along_axis(lambda v: v - max(v), 1, X @ w.T + b)
    probs = np.exp(z)
    denom = np.sum(probs, axis=1)
    s = np.zeros((N, C))
    for i in range(N):
        s[i] = probs[i] / denom[i]
    for i in range(N):
        s[i][y[i]] += -1
    for c in range(C):
        grad_w = s.T[c] @ X
        w[c] += (step_size / N) * grad_w
        b[c] += (step_size / N) * sum(s.T[c])

问题 #

1. 第二段代码为何无法得到正确结果？应如何修正？
2. 更重要的是，如何进一步优化这段代码，以充分利用Numpy的向量化运算？

解答 #

问题1：错误原因与修正

你的优化代码核心问题是梯度更新的符号搞反了，同时还有一些可以简化的低效操作。

我们对比初始代码和优化代码的更新逻辑：

• 初始代码中，权重更新是 w[i] -= (step_size/N) * term * X[j]，其中term对类别匹配标签时是(probs[j][i]/denom[j])-1，否则是probs[j][i]/denom[j]
• 你构建的s矩阵中，s[j][c]正好等于这个term（标签位置减1后），所以梯度应该是sum(term * X[j])，而初始代码是减去这个梯度乘以步长/N
• 但你的优化代码里写的是w[c] += (step_size/N) * grad_w，这就相当于做了反向更新，导致模型收敛方向错误

另外，你的代码里还有两个可以优化的点：

• 手动循环赋值s[i] = probs[i]/denom[i]完全没必要，Numpy的广播可以直接处理
• 用Python内置sum替代Numpy的np.sum会降低效率

修正后的代码如下：

for _ in range(max_iterations):
    z = np.apply_along_axis(lambda v: v - max(v), 1, X @ w.T + b)
    probs = np.exp(z)
    # 加keepdims=True让denom保持(N,1)形状，支持广播除法
    denom = np.sum(probs, axis=1, keepdims=True)
    # 直接广播除法，替代循环赋值
    s = probs / denom
    # 用向量化索引给标签位置减1，替代循环
    s[np.arange(N), y] -= 1
    # 计算梯度
    grad_w = s.T @ X
    grad_b = np.sum(s, axis=0)
    # 修正符号：改为减去梯度乘以步长/N
    w -= (step_size / N) * grad_w
    b -= (step_size / N) * grad_b

问题2：充分利用Numpy向量化的终极优化

要彻底发挥Numpy的性能，我们需要去掉所有显式循环，包括np.apply_along_axis（它本质上是Python循环的包装），全部用Numpy的向量化操作替代：

for _ in range(max_iterations):
    # 计算logits，形状(N, C)
    logits = X @ w.T + b
    # 数值稳定：每行减去最大值，避免指数爆炸，用向量化操作替代apply_along_axis
    logits -= np.max(logits, axis=1, keepdims=True)
    # 计算概率
    probs = np.exp(logits)
    probs /= np.sum(probs, axis=1, keepdims=True)
    # 调整概率矩阵：标签位置减1，得到梯度项
    probs[np.arange(N), y] -= 1
    # 计算梯度（直接除以N，减少一次乘法）
    grad_w = probs.T @ X / N
    grad_b = np.sum(probs, axis=0) / N
    # 更新权重和偏置
    w -= step_size * grad_w
    b -= step_size * grad_b

关键优化细节：

1. 替换np.apply_along_axis：用np.max(logits, axis=1, keepdims=True)得到每行最大值，然后通过广播直接减去，比apply_along_axis快数倍
2. 全向量化操作：所有除法、减法、索引都利用Numpy的广播机制，彻底去掉所有Python循环
3. 简化梯度计算：直接在梯度计算时除以N，减少一次额外的乘法操作
4. 高效索引：用np.arange(N)配合y实现向量化的标签位置修改，避免逐个样本循环

这样的代码不仅速度快，而且更简洁易读，完全利用了Numpy的底层优化（比如C语言实现的矩阵运算）。

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