Logistic Regression

原始题目：LeetGPU - Logistic Regression

题目描述

在 GPU 上求解逻辑回归问题。给定特征矩阵 $X$（$n\_samples \times n\_features$）和二元目标向量 $y$（长度 $n\_samples$，仅含 0 和 1），计算最大化对数似然的系数向量 $\beta$：

$$\max_{\beta} \sum_{i=1}^{n} \left[y_i \log(p_i) + (1 - y_i) \log(1 - p_i)\right]$$

其中 $p_i = \sigma(X_i^T \beta)$，$\sigma(z) = \frac{1}{1 + e^{-z}}$ 为 sigmoid 函数。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终系数必须存储在 beta 向量中。
• 目标向量 $y$ 仅包含 0 和 1。

示例

Input:  X (8 samples × 2 features):
          [[2,1],[1,2],[3,3],[1.5,2.5],[-1,-2],[-2,-1],[-1.5,-2.5],[-3,-3]]
        y: [1,1,1,0,0,0,1,0]
Output: β: [2.26, -1.29]

约束条件

• $1 \le n\_samples \le 100{,}000$，$1 \le n\_features \le 1{,}000$。
• $n\_samples \ge n\_features$。
• $-10.0 \le X$ 中的值 $\le 10.0$。
• 绝对容差 $10^{-2}$，相对容差 $10^{-2}$。
• 性能测试在 $n\_features = 8,\ n\_samples = 16$ 的规模下进行。

解题思路

逻辑回归没有闭式解，需要通过迭代优化求解（如 Newton-Raphson、梯度下降或 IRLS）。每次迭代需要计算梯度 $X^T(p - y)$ 和 Hessian $X^T W X$（$W$ 是对角权重矩阵 $p_i(1-p_i)$）。计算瓶颈在矩阵乘法（$X^T W X$），这是一个对称 rank-k 更新，可以利用 GEMM 加速。当 $n\_features$ 较小时，多项迭代在 GPU 上的 overhead 可能大于计算收益。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
// IRLS for logistic regression, simplified
__global__ void lr_grad(const float* X, const float* y, const float* beta, float* grad, int N, int F) {
    int f=threadIdx.x+blockIdx.x*blockDim.x;
    if(f<F){float sum=0;for(int i=0;i<N;i++){float z=0;for(int j=0;j<F;j++)z+=X[i*F+j]*beta[j];sum+=(1.0f/(1.0f+expf(-z))-y[i])*X[i*F+f];}grad[f]=sum/N;}
}
extern "C" void solve(const float* X, const float* y, float* beta, int N, int F) {
    // Simplified: real impl uses Newton-Raphson iterations
    float *grad; cudaMalloc(&grad,F*sizeof(float));
    lr_grad<<<(F+255)/256,256>>>(X,y,beta,grad,N,F); cudaDeviceSynchronize();
    cudaFree(grad);
}

Triton

import torch
def solve(X, y, n_iter=100, lr=0.1):
    beta = torch.zeros(X.shape[1], device=X.device)
    for _ in range(n_iter):
        p = torch.sigmoid(X @ beta)
        grad = X.T @ (p - y) / len(y)
        beta -= lr * grad
    return beta