arXiv来源发表:
核方法求解器不再被预条件子内存卡住:分块草图让模型参数扩到千万级,精度还更高
导语
核岭回归的预条件子内存不再随模型参数增长,Bloks 把参数规模推到千万级,在 Higgs 等大规模基准上取得与当前方法相当或更好的精度。
正文
核岭回归的求解器现在可以把模型参数规模与内存开销分开:预条件子只保留一个块,内存不再随数据量和参数个数增长。 此前基于草图的预条件子需要存储一个与参数个数平方相关的矩阵,参数一大就内存溢出,只能把模型压小。 在 Higgs 数据集上,Falkon 因预条件子的二次存储需求无法越过某一参数规模而内存溢出,Bloks 在保持内存不变的情况下把参数扩到千万级,测试精度随之明显提升。
做法是把 Nyström 中心切成互不重叠的小块,对每块单独构造一个草图后的对角块,再把这些块拼成预条件子。 标准分块 Jacobi 预条件子虽然内存不随参数增长,但每块都要算一次 Gram 乘积,数值不稳定且代价过高。 每个草图块写成 Cholesky 因子的形式,用两次原地 Cholesky 分解只存一个矩阵,块之间解耦,可顺序处理也可并行;实验在单张 A100 上完成,训练时间与主流方法相当。
接下来
报告覆盖的是 Nyström 草图核岭回归上的求解器,实验在单张 A100 上以单精度运行,对比对象为 EigenPro 4、Joker、ASkotch 和 Falkon。作者提出后续可把方法扩展到其他损失函数、研究不同的分块策略,以及加入捕捉块间相互作用的粗校正。
作者指出,分块近似带来的额外正则化使 Bloks 在训练误差上不如标准分块 Jacobi,但测试表现几乎一致,这一现象的理论分析留待未来工作。块大小增大时训练时间超线性增长,且当所有块无法同时驻留内存时预条件子代价不再可忽略。
