[论文解读] Auto-Differentiating Linear Algebra

Matthias Seeger, Asmus Hetzel+3 more|arXiv (Cornell University)|Oct 24, 2017

Gaussian Processes and Bayesian Inference参考文献 18被引用 17

一句话总结

本文将可微分的线性代数算子——Cholesky分解、LQ分解、对称特征值分解和SVD——引入MXNet深度学习框架，实现了端到端训练混合模型，结合深度神经网络与贝叶斯方法（如高斯过程和卡尔曼滤波器）。通过利用高度优化的原地BLAS/LAPACK调用并最小化内存开销，该实现使CPU性能提升三倍以上，GPU性能持续优于TensorFlow。

ABSTRACT

Development systems for deep learning (DL), such as Theano, Torch, TensorFlow, or MXNet, are easy-to-use tools for creating complex neural network models. Since gradient computations are automatically baked in, and execution is mapped to high performance hardware, these models can be trained end-to-end on large amounts of data. However, it is currently not easy to implement many basic machine learning primitives in these systems (such as Gaussian processes, least squares estimation, principal components analysis, Kalman smoothing), mainly because they lack efficient support of linear algebra primitives as differentiable operators. We detail how a number of matrix decompositions (Cholesky, LQ, symmetric eigen) can be implemented as differentiable operators. We have implemented these primitives in MXNet, running on CPU and GPU in single and double precision. We sketch use cases of these new operators, learning Gaussian process and Bayesian linear regression models, where we demonstrate very substantial reductions in implementation complexity and running time compared to previous codes. Our MXNet extension allows end-to-end learning of hybrid models, which combine deep neural networks (DNNs) with Bayesian concepts, with applications in advanced Gaussian process models, scalable Bayesian optimization, and Bayesian active learning.

研究动机与目标

• 在深度学习系统中实现关键线性代数原语的高效自动微分。
• 降低经典机器学习模型（如高斯过程和贝叶斯线性回归）的实现复杂度与运行时开销。
• 支持混合深度学习与贝叶斯模型，包括变分自编码器和深度高斯过程。
• 通过原地计算和直接调用CPU与GPU上经过调优的BLAS/LAPACK库，优化性能。
• 将深度学习框架的能力扩展至标准神经网络层之外，涵盖基础机器学习算法。

提出的方法

• 在MXNet中将Cholesky分解、LQ分解、对称特征值分解和SVD分解实现为一等可微分算子。
• 使用低级C++接口直接调用高度优化的BLAS和LAPACK例程，以实现高性能计算。
• 设计原地操作，以最小化临时内存分配，提升内存效率。
• 支持CPU和GPU上的单精度与双精度浮点数运算。
• 将算子集成到MXNet的符号式与命令式API（Gluon）中，支持灵活的模型构建。
• 推导并实现每种算子的反向模式微分规则，完整反向传播推导详见附录。

实验结果

研究问题

• RQ1 能否在深度学习框架中实现Cholesky和SVD等矩阵分解的可微分且高效执行？
• RQ2 将线性代数原语作为一等算子暴露后，能在多大程度上提升性能并降低实现复杂度？
• RQ3 使用这些算子在多大程度上可简化并加速结合深度神经网络与贝叶斯方法的混合模型？
• RQ4 所提出的MXNet算子在性能特征上与TensorFlow和PyTorch中的现有实现相比如何？
• RQ5 这些算子能否在极少代码修改下支持新型可扩展模型（如随机变分深度高斯过程）？

主要发现

• MXNet中可微分线性代数算子的实现，在CPU（Intel MKL）上相比等效的TensorFlow操作，执行速度提升三倍以上。
• 在所有测试工作负载中，该算子在GPU上持续优于TensorFlow，且由于采用原地计算，内存使用更少。
• 通过极少代码修改实现了新模型——随机变分深度高斯过程（SV-DGP），在MNIST数据集上的测试对数似然值高于VAE-DGP（5-20-50架构下分别为724.83 vs. 723.65）。
• 该实现使复杂混合模型（如结合变分自编码器的深度高斯过程）的端到端训练成为可能，且代码复杂度显著降低。
• 使用高度优化的原地BLAS和LAPACK调用，使内存开销极低，使该算子适用于高斯过程和卡尔曼平滑等内存密集型应用。
• 该框架通过MXNet的Gluon API支持静态与动态计算图，可灵活实验混合模型。