[论文解读] Inv-ASKIT: A Parallel Fast Diret Solver for Kernel Matrices

Chenhan D. Yu, William B. March+2 more|arXiv (Cornell University)|Feb 3, 2016

Stochastic Gradient Optimization Techniques参考文献 19被引用 4

一句话总结

Inv-ASKIT 是一种用于核矩阵的并行快速直接求解器，通过基于 ASKIT 的矩阵近似利用分层低秩结构，实现了 $\mathcal{O}(N\log N)$ 的线性系统 $\lambda I + K$ 求解。其分解过程耗时 $\mathcal{O}(N\log^2 N)$，可扩展至 4,096 个核心，在高维数据上求解 32M×32M 矩阵时达到峰值 FLOPS 的 50%，相比以往方法实现了多数量级的性能提升。

ABSTRACT

We present a parallel algorithm for computing the approximate factorization of an $N$-by-$N$ kernel matrix. Once this factorization has been constructed (with $N \log^2 N $ work), we can solve linear systems with this matrix with $N \log N $ work. Kernel matrices represent pairwise interactions of points in metric spaces. They appear in machine learning, approximation theory, and computational physics. Kernel matrices are typically dense (matrix multiplication scales quadratically with $N$) and ill-conditioned (solves can require 100s of Krylov iterations). Thus, fast algorithms for matrix multiplication and factorization are critical for scalability. Recently we introduced ASKIT, a new method for approximating a kernel matrix that resembles N-body methods. Here we introduce INV-ASKIT, a factorization scheme based on ASKIT. We describe the new method, derive complexity estimates, and conduct an empirical study of its accuracy and scalability. We report results on real-world datasets including "COVTYPE" ($0.5$M points in 54 dimensions), "SUSY" ($4.5$M points in 8 dimensions) and "MNIST" (2M points in 784 dimensions) using shared and distributed memory parallelism. In our largest run we approximately factorize a dense matrix of size 32M $ imes$ 32M (generated from points in 64 dimensions) on 4,096 Sandy-Bridge cores. To our knowledge these results improve the state of the art by several orders of magnitude.

研究动机与目标

• 为解决高维场景下稠密核矩阵的可扩展性与病态问题，其中标准直接求解器与迭代求解器成本过高。
• 开发一种并行直接求解器，实现核矩阵在大规模高维数据集上的 $\mathcal{O}(N\log N)$ 求解时间与 $\mathcal{O}(N\log^2 N)$ 分解成本。
• 支持共享内存与分布式内存并行计算，实现核矩阵分解与求逆，适用于广泛的核函数，包括变带宽与非对称情形。
• 在真实数据集（如 COVTYPE、SUSY 和 MNIST）上展示高精度与强可扩展性，即使在极端规模（如 32M×32M 矩阵）下亦表现优异。
• 为 GMRES 等迭代求解器提供预条件子，显著加速核基机器学习应用中的收敛速度。

提出的方法

• 该方法使用 ASKIT 将核矩阵 $K$ 分解为块对角矩阵、低秩矩阵与稀疏矩阵 ($S$)，实现分层低秩近似。
• 通过递归应用 Sherman-Morrison-Woodbury 公式，分别对块对角与低秩分量求逆，从而实现 $\lambda I + K$ 的高效求逆，降低整体计算复杂度。
• 利用基于空间邻近性与核函数衰减特性的分层树结构，以 $\mathcal{O}(N\log^2 N)$ 时间完成分解。
• 算法支持共享内存与分布式内存并行计算，采用动态负载均衡与通信避免型数据布局，有效降低通信开销。
• 在迭代求解中，Inv-ASKIT 作为预条件子使用，将 $\tilde{K}$ 的逆作为初始猜测，显著加速 GMRES 的收敛。
• 该方法每项核函数计算仅需 $\mathcal{O}(d)$ 时间，适用于高维输入空间。

实验结果

研究问题

• RQ1 能否在高维数据上实现 $\mathcal{O}(N\log N)$ 求解时间的核矩阵直接求解器，同时保持高精度？
• RQ2 分层低秩矩阵近似是否能实现稠密核矩阵在共享与分布式内存系统上的可扩展、并行分解与求逆？
• RQ3 作为预条件子用于核岭回归时，Inv-ASKIT 与 Krylov 迭代方法相比，在精度与收敛速度方面表现如何？
• RQ4 Inv-ASKIT 在极端规模数据集（如 32M×32M 矩阵）与大规模核心数（如 4,096 核）下的实际可扩展性如何？
• RQ5 该方法是否能处理变带宽与非对称核函数，而无需依赖对称性假设？

主要发现

• Inv-ASKIT 在 4,096 个核心上成功分解了 32M×32M 的核矩阵，达到峰值 FLOPS 的 50%，相比以往最先进方法实现了多数量级的性能提升。
• 在 COVTYPE 数据集（54维，0.5M 个点）上，使用 $\tilde{K}$ 时分类准确率达 96%，加入最近邻稀疏校正 $S$ 后提升至 97%。
• 在 SUSY 数据集（8维，450万个点）上，准确率达到 79%；在 MNIST 数据集（784维，200万个点）上，核回归任务准确率达 100%。
• 作为 GMRES 的预条件子时，Inv-ASKIT 将迭代次数从 100 次减少至 0 次（立即收敛），证明其高精度与高效预条件能力。
• 在弱扩展性实验中，Inv-ASKIT 在 4,096 个核心上保持了高效率，通信与计算成本均呈有利扩展趋势。
• 相比非预条件化的 GMRES，该方法实现了 2 倍至 3 倍的加速，并在合成与真实数据集上均表现出强可扩展性。