A Tensor Algebra Compiler for Sparse Differentiation

Sparse tensors are prevalent in many data-intensive applications. However, existing automatic differentiation (AD) frameworks are tailored towards dense tensors, which makes it a challenge to efficiently compute gradients through sparse tensor operations. This is due to irregular sparsity patterns t...

Full description

Saved in:
Bibliographic Details
Published in:Proceedings / International Symposium on Code Generation and Optimization pp. 1 - 12
Main Authors: Shaikhha, Amir, Huot, Mathieu, Hashemian, Shideh
Format: Conference Proceeding
Language:English
Published: IEEE 02.03.2024
Subjects:
Automatic Differentiation
Codes
Dictionaries
DSL
Optimization
Pipelines
Scalability
Semi-Ring Dictionaries
Sparse Tensor Algebra
Tensors
ISSN:2643-2838
Online Access:Get full text
Tags: Add Tag
No Tags, Be the first to tag this record!
Description
Summary:Sparse tensors are prevalent in many data-intensive applications. However, existing automatic differentiation (AD) frameworks are tailored towards dense tensors, which makes it a challenge to efficiently compute gradients through sparse tensor operations. This is due to irregular sparsity patterns that can result in substantial memory and computational overheads. We propose a novel framework that enables the efficient AD of sparse tensors. The key aspects of our work include a compilation pipeline leveraging two intermediate DSLs with AD-agnostic domain-specific optimizations followed by efficient C++ code generation. We showcase the effectiveness of our framework in terms of performance and scalability through extensive experimentation, outperforming state-of-the-art alternatives across a variety of synthetic and real-world datasets.
ISSN:2643-2838
DOI:10.1109/CGO57630.2024.10444787