Scikit-learn QuantileRegressor内存分配错误问题求助

问题现象 #

• 使用scikit-learn的QuantileRegressor拟合分位数回归模型时，触发内存分配错误，报错信息如下：

numpy.core._exceptions._ArrayMemoryError: Unable to allocate 55.9 GiB for an array with shape (86636, 86636) and data type float64

• 相同输入数据（特征矩阵形状(86636, 4)，目标数组形状(86636, 1)）在statsmodels的QuantReg中可正常运行，耗时仅2-3秒；若大幅缩减sklearn版本的输入数据至1000行以内，仍需数分钟才能完成运行。

复现代码 #

import pandas as pd
import statsmodels.api as sm
from sklearn.linear_model import QuantileRegressor

training_df = pd.read_csv("/path/to/training_df.csv") # 86,000行

FEATURES = [
    "feature_1",
    "feature_2",
    "feature_3",
    "feature_4",
]

TARGET = "target"

# statsmodels版本正常运行
model_statsmodels = sm.QuantReg(training_df[TARGET], training_df[FEATURES]).fit(q=0.5)

# sklearn版本触发内存错误
model_sklearn = QuantileRegressor(quantile=0.5, alpha=0)
model_sklearn.fit(training_df[FEATURES], training_df[TARGET])

原因分析 #

scikit-learn的QuantileRegressor默认采用内点法作为优化器，当设置alpha=0（无正则化）且数据量较大时，该算法需要构建一个(n_samples, n_samples)规模的矩阵，对于86636条样本来说，这个矩阵需要约55.9GiB的内存，远超普通机器的内存上限。

而statsmodels的QuantReg默认使用单纯形法，该算法针对分位数回归的优化问题做了专门设计，无需构建大规模矩阵，因此在大数据量下内存占用更低、运行效率更高。

解决思路 #

• 切换优化算法：给scikit-learn的QuantileRegressor指定solver='highs'（需确保scipy>=1.6.0已安装），该求解器基于单纯形法实现，内存占用和运行效率会接近statsmodels的版本：

  model_sklearn = QuantileRegressor(quantile=0.5, alpha=0, solver='highs')
  model_sklearn.fit(training_df[FEATURES], training_df[TARGET])
  

• 添加正则化：若必须使用内点法，可设置alpha>0（如alpha=1e-3），正则化会大幅降低内点法的内存消耗，但会引入轻微的正则化偏差，需根据业务场景权衡。
• 继续使用statsmodels：如果不需要scikit-learn的流水线、交叉验证等生态特性，直接沿用statsmodels的QuantReg是最省心的方案，性能表现更适配当前数据规模。

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