用KFAC提升双层优化的超梯度计算效率与精度
Efficient Bilevel Optimization with KFAC-Based Hypergradients
- 基于隐函数定理，引入KFAC近似曲率信息
- 在BERT级模型上实现更高精度的超梯度，内存开销小
- 适用于元学习和AI安全等复杂任务
双层优化（BO）广泛应用于机器学习，但其扩展性受限于超梯度计算，需求解逆海森向量积（IHVP）。现有方法常使用简化的近似，如单步梯度回滚或恒等/短奈曼展开，忽略曲率信息。本文基于隐函数定理，引入克罗内克分解近似曲率（KFAC），获得更精准的超梯度，在性能与效率之间取得更好平衡，优于共轭梯度（CG）或奈曼方法，并持续超越回滚法。我们在元学习与AI安全等多种任务中验证该方法，对高达BERT规模的模型，证明曲率信息在大规模下依然有价值，且KFAC仅带来适度的内存与运行时间开销。代码已开源：https://github.com/liaodisen/NeuralBo。
原文摘要 · Abstract (English)
Bilevel optimization (BO) is widely applicable to many machine learning problems. Scaling BO, however, requires repeatedly computing hypergradients, which involves solving inverse Hessian-vector products (IHVPs). In practice, these operations are often approximated using crude surrogates such as one-step gradient unrolling or identity/short Neumann expansions, which discard curvature information. We build on implicit function theorem-based algorithms and propose to incorporate Kronecker-factored approximate curvature (KFAC), yielding curvature-aware hypergradients with a better performance efficiency trade-off than Conjugate Gradient (CG) or Neumann methods and consistently outperforming unrolling. We evaluate this approach across diverse tasks, including meta-learning and AI safety problems. On models up to BERT, we show that curvature information is valuable at scale, and KFAC can provide it with only modest memory and runtime overhead. Our implementation is available at https://github.com/liaodisen/NeuralBo.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。