用卡尔曼滤波思想高效建模梯度不确定性,提升训练效果
KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products
- 通过递归更新梯度协方差积,直接估计参数协方差矩阵
- 在图像分类与语言建模上达到或超越顶尖优化器性能
- 无需存储完整协方差矩阵,兼顾效率与高阶信息捕捉
我们提出KOALA++,一种可扩展的基于卡尔曼滤波的优化算法,显式建模神经网络训练中的结构化梯度不确定性。与依赖昂贵二阶梯度计算的二阶方法不同,该方法通过递归更新紧凑的梯度协方差积,直接估计参数协方差矩阵。这一设计改进了原始KOALA框架(假设协方差对角),在不存储完整协方差矩阵且避免大规模矩阵求逆的前提下,隐式捕捉更丰富的不确定性结构。在图像分类和语言建模等多样化任务中,KOALA++的准确率与当前最优的一阶和二阶优化器相当或更优,同时保持了一阶方法的计算效率。
原文摘要 · Abstract (English)
We propose KOALA++, a scalable Kalman-based optimization algorithm that explicitly models structured gradient uncertainty in neural network training. Unlike second-order methods, which rely on expensive second order gradient calculation, our method directly estimates the parameter covariance matrix by recursively updating compact gradient covariance products. This design improves upon the original KOALA framework that assumed diagonal covariance by implicitly capturing richer uncertainty structure without storing the full covariance matrix and avoiding large matrix inversions. Across diverse tasks, including image classification and language modeling, KOALA++ achieves accuracy on par or better than state-of-the-art first- and second-order optimizers while maintaining the efficiency of first-order methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。