[论文解读] Optimal Sketching for Trace Estimation

Shuli Jiang, Hai Pham+3 more|arXiv (Cornell University)|Nov 1, 2021

Stochastic Gradient Optimization Techniques参考文献 33被引用 5

一句话总结

本文提出了一种用于迹估计的最优非自适应压缩算法，其查询复杂度为 $O(\sqrt{\log(1/\delta)}/\epsilon + \log(1\delta))$，可在失败概率 $\delta$ 下实现 $(1\pm\epsilon)$-近似迹估计，其复杂度与理论下界相比仅相差 $\log\log(1/\delta)$ 因子。该工作通过证明非自适应算法可同时实现对 $\epsilon$ 和 $\delta$ 的最优依赖关系，从而弥合了自适应与非自适应方法之间的差距。

ABSTRACT

Matrix trace estimation is ubiquitous in machine learning applications and has traditionally relied on Hutchinson's method, which requires $O(\log(1/δ)/ε^2)$ matrix-vector product queries to achieve a $(1 \pm ε)$-multiplicative approximation to $ ext{tr}(A)$ with failure probability $δ$ on positive-semidefinite input matrices $A$. Recently, the Hutch++ algorithm was proposed, which reduces the number of matrix-vector queries from $O(1/ε^2)$ to the optimal $O(1/ε)$, and the algorithm succeeds with constant probability. However, in the high probability setting, the non-adaptive Hutch++ algorithm suffers an extra $O(\sqrt{\log(1/δ)})$ multiplicative factor in its query complexity. Non-adaptive methods are important, as they correspond to sketching algorithms, which are mergeable, highly parallelizable, and provide low-memory streaming algorithms as well as low-communication distributed protocols. In this work, we close the gap between non-adaptive and adaptive algorithms, showing that even non-adaptive algorithms can achieve $O(\sqrt{\log(1/δ)}/ε+ \log(1/δ))$ matrix-vector products. In addition, we prove matching lower bounds demonstrating that, up to a $\log \log(1/δ)$ factor, no further improvement in the dependence on $δ$ or $ε$ is possible by any non-adaptive algorithm. Finally, our experiments demonstrate the superior performance of our sketch over the adaptive Hutch++ algorithm, which is less parallelizable, as well as over the non-adaptive Hutchinson's method.

研究动机与目标

• 解决迹估计中自适应与非自适应算法之间的性能差距，尤其是在高概率设置下。
• 通过实现对 $\epsilon$ 和 $\delta$ 的最优依赖关系，弥合非自适应压缩算法在查询复杂度上的理论差距。
• 证明非自适应方法可与 Hutch++ 等自适应算法在效率上相当，同时保有并行性和低内存流式处理的优势。
• 建立紧致的下界，表明任何非自适应算法在所推导的查询复杂度上都无法进一步改进，最多相差 $\log\log(1/\delta)$ 因子。
• 通过实验验证，所提出的压缩方法在实践中优于非自适应的 Hutchinson 方法和自适应的 Hutch++ 方法，尤其在分布式和流式环境中表现更优。

提出的方法

• 设计一种新颖的非自适应压缩算法，结合随机矩阵-向量查询与方差减少技术，以实现最优查询复杂度。
• 采用与维度无关的查询向量分布，确保在不同矩阵维度下均具有鲁棒性和可扩展性。
• 应用先进的集中不等式与 KL 散度分析，界定向量分布之间的总变差距离，证明在特定条件下不可区分性。
• 利用 Pinsker 不等式与条件 KL 散度推导出区分算法失败概率的下界，从而得出查询复杂度的下界。
• 引入 NA-Hutch++ 算法的改进版本并提供更优的分析，证明其在高概率下可实现所声称的查询复杂度。
• 通过高斯混合模型构造一个困难实例，并利用信息论论证证明，任何非自适应算法都必须进行 $\Omega(\sqrt{\log(1/\delta)}/\epsilon + \log(1/\delta))$ 次查询，才能以 $1 - \delta$ 的概率成功。

实验结果

研究问题

• RQ1 非自适应压缩算法能否在迹估计中实现与 Hutch++ 等自适应算法相同的查询复杂度？
• RQ2 非自适应迹估计算法的查询复杂度对 $\epsilon$ 和 $\delta$ 的最优依赖关系是什么？
• RQ3 是否存在非自适应方法在查询复杂度上可进一步减少的根本限制？若有，其上限是什么？
• RQ4 所提出的算法能否在实践中优于非自适应的 Hutchinson 方法和自适应的 Hutch++ 方法，尤其是在高并行或流式环境中？
• RQ5 非自适应迹估计的查询复杂度理论下界有多紧？是否与上界仅相差对数因子？

主要发现

• 所提出的非自适应压缩算法实现了 $O(\sqrt{\log(1/\delta)}/\epsilon + \log(1/\delta))$ 的查询复杂度，与理论下界相比仅相差 $\log\log(1/\delta)$ 因子。
• 该算法通过实现对 $\epsilon$ 和 $\delta$ 的最优依赖关系，弥合了非自适应与自适应方法之间的差距，解决了长期存在的开放问题。
• 证明了匹配的下界：任何非自适应算法都无法实现优于 $\Omega(\sqrt{\log(1/\delta)}/\epsilon + \log(1/\delta))$ 的查询复杂度，最多相差 $\log\log(1/\delta)$ 因子。
• 实验结果表明，所提出的压缩方法在收敛速度和可并行性方面均优于非自适应的 Hutchinson 方法和自适应的 Hutch++ 算法。
• 该算法在保持强理论保证的同时，具有高度可并行性，适用于低内存流式处理和分布式计算环境。
• 分析表明，任何非自适应算法通过区分两个分布的失败概率均被下界限制为 $\delta$，从而在某些场景下得出查询复杂度的下界为 $\Omega(\log(1/\delta)/\log\log(1/\delta))$。