CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra

要約

機械学習と科学の多くの分野には、固有分解、線形システムの解法、行列指数の計算、トレース推定などの大規模な線形代数問題が含まれます。
関連する行列には、クロネッカー構造、畳み込み構造、ブロック対角構造、和構造、または積構造が含まれることがよくあります。
この論文では、CoLA (Compositional Linear Algebra) と呼ばれる、機械学習における大規模線形代数問題のための単純だが一般的なフレームワークを提案します。
線形演算子の抽象化と構成的なディスパッチ ルールを組み合わせることで、CoLA はメモリとランタイムに効率的な数値アルゴリズムを自動的に構築します。
さらに、CoLA は、JAX と PyTorch の両方でメモリ効率の高い自動微分、低精度計算、GPU アクセラレーションを提供すると同時に、複数のディスパッチを介してダウンストリーム パッケージで新しいオブジェクト、操作、ルールにも対応します。
CoLA は、多くの代数演算を高速化できると同時に、行列構造とアルゴリズムのプロトタイプ作成を容易にし、線形代数を必要とするほぼすべての計算作業に魅力的なドロップイン ツールを提供します。
偏微分方程式、ガウス過程、等変モデル構築、教師なし学習など、幅広いアプリケーションにわたるその有効性を紹介します。

要約(オリジナル)

Many areas of machine learning and science involve large linear algebra problems, such as eigendecompositions, solving linear systems, computing matrix exponentials, and trace estimation. The matrices involved often have Kronecker, convolutional, block diagonal, sum, or product structure. In this paper, we propose a simple but general framework for large-scale linear algebra problems in machine learning, named CoLA (Compositional Linear Algebra). By combining a linear operator abstraction with compositional dispatch rules, CoLA automatically constructs memory and runtime efficient numerical algorithms. Moreover, CoLA provides memory efficient automatic differentiation, low precision computation, and GPU acceleration in both JAX and PyTorch, while also accommodating new objects, operations, and rules in downstream packages via multiple dispatch. CoLA can accelerate many algebraic operations, while making it easy to prototype matrix structures and algorithms, providing an appealing drop-in tool for virtually any computational effort that requires linear algebra. We showcase its efficacy across a broad range of applications, including partial differential equations, Gaussian processes, equivariant model construction, and unsupervised learning.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google