[论文解读] Differentiating a Tensor Language

Gilbert Bernstein, Michael W. Mara+3 more|arXiv (Cornell University)|Aug 25, 2020

Parallel Computing and Optimization Techniques参考文献 29被引用 5

一句话总结

本文提出了一种纯函数式、可证明高效的张量程序反向模式导数计算方法，通过使用Iverson的指示函数显式建模稀疏性。通过利用一种新的“张量SSA”规范形式以及基于内积的自动微分推导，该方法在不使用可变性的前提下实现了廉价梯度原则，确保导数代码的渐近效率与原函数相当。

ABSTRACT

How does one compile derivatives of tensor programs, such that the resulting code is purely functional (hence easier to optimize and parallelize) and provably efficient relative to the original program? We show that naively differentiating tensor code---as done in popular systems like Tensorflow and PyTorch---can cause asymptotic slowdowns in pathological cases, violating the Cheap Gradients Principle. However, all existing automatic differentiation methods that guarantee this principle (for variable size data) do so by relying on += mutation through aliases/pointers---which complicates downstream optimization. We provide the first purely functional, provably efficient, adjoint/reverse-mode derivatives of array/tensor code by explicitly accounting for sparsity. We do this by focusing on the indicator function from Iverson's APL. We also introduce a new "Tensor SSA" normal form and a new derivation of reverse-mode automatic differentiation based on the universal property of inner-products.

研究动机与目标

• 为解决张量程序中朴素自动微分的低效问题，其可能导致梯度计算的渐近性能下降。
• 通过避免诸如 += 变更等命令式副作用，保留数据并行性并支持后续优化。
• 在密集张量和稀疏张量的纯函数式环境中，形式化保证廉价梯度原则。
• 引入一种新的规范形式“张量SSA”，以实现高效且正确的导数代码生成。
• 为基于内积通用性质的反向模式AD提供数学上严谨的基础。

提出的方法

• 该方法使用Iverson的指示函数表示布尔谓词，并在张量运算中显式建模稀疏性。
• 提出一种新的“张量SSA”规范形式，支持张量程序的静态分析与优化。
• 利用内积的通用性质推导导数，确保正确性与高效性。
• 在伴随计算中考虑稀疏性，避免不必要的内存写入。
• 应用代数简化技术，即使在病态情况下也能保持渐近效率。
• 通过将稀疏性视为一等公民，避免 += 变更，实现纯函数式编译。

实验结果

研究问题

• RQ1 我们能否在不依赖命令式副作用的纯函数式张量语言中实现廉价梯度原则？
• RQ2 如何形式化建模张量运算中的稀疏性，以确保高效的导数计算？
• RQ3 何种规范形式能够支持数据并行张量程序中高效且正确的反向模式自动微分？
• RQ4 内积的通用性质能否用于推导出一个可证明高效的AD系统？
• RQ5 如何系统性地避免现有框架（如PyTorch、TensorFlow）中的渐近低效问题？

主要发现

• 所提方法在纯函数式环境中实现了廉价梯度原则，确保导数计算的性能至多比原程序慢一个常数因子。
• 通过使用Iverson指示函数建模稀疏性，该方法避免了伴随计算中的不必要的内存写入，例如在diag和trace运算中。
• 该方法证明了前向传递与反向传递之间的成本等价性，即使在考虑内积和求和调整后依然成立。
• 张量SSA规范形式支持静态分析与优化，保留了数据并行性，并支持高效编译。
• 在病态情况下，该方法优于朴素微分——例如，当优化时，trace(A) + ... + trace(A) 的梯度不再为 O(kN²)，而是 O(kN)。
• 该方法避免了对 += 变更的依赖，从而保持了函数纯度，并支持更强的编译器优化。