Scorch: A Library for Sparse Deep Learning

Abstract

The rapid growth in the size of deep learning models strains the capabilities of traditional dense computation paradigms. Leveraging sparse computation has become increasingly popular for training and deploying large-scale models, but existing deep learning frameworks lack extensive support for sparse operations. To bridge this gap, we introduce Scorch, a library that seamlessly integrates efficient sparse tensor computation into the PyTorch ecosystem, with an initial focus on inference workloads on CPUs. Scorch provides a flexible and intuitive interface for sparse tensors, supporting diverse sparse data structures. Scorch introduces a compiler stack that automates key optimizations, including automatic loop ordering, tiling, and format inference. Combined with a runtime that adapts its execution to both dense and sparse data, Scorch delivers substantial speedups over hand-written PyTorch Sparse (torch.sparse) operations without sacrificing usability. More importantly, Scorch enables efficient computation of complex sparse operations that lack hand-optimized PyTorch implementations. This flexibility is crucial for exploring novel sparse architectures. We demonstrate Scorch’s ease of use and performance gains on diverse deep learning models across multiple domains. With only minimal code changes, Scorch achieves 1.05–5.78$\times$ speedups over PyTorch Sparse on end-to-end tasks. Scorch’s seamless integration and performance gains make it a valuable addition to the PyTorch ecosystem. We believe Scorch will enable wider exploration of sparsity as a tool for scaling deep learning and inform the development of other sparse libraries.

1 Introduction

The rapid increase in model size and sophistication has driven the success of deep learning across domains like computer vision, natural language processing, and speech recognition. This growth strains the capabilities of traditional dense computation paradigms, as model parameters often exceed available processing resources and memory. Sparse computation techniques offer a promising solution by focusing compute and storage only on the most relevant elements [25].

Sparsity arises naturally in many contexts within deep learning and can generally be categorized into three groups: data sparsity, weight sparsity, and activation sparsity. It arises in data, as many emerging application domains involve inherently sparse data representations. For instance, real-world graphs often exhibit power-law distributions in node connectivity, yielding sparse adjacency matrices [59]. Similarly, the high-dimensional feature spaces used in recommender systems [46] and natural language processing [12] are highly sparse.

Sparsity also arises within models, where weight and activation sparsity can provide significant efficiency benefits. Techniques like pruning [24, 8, 41] have shown that weights can be sparsified with little impact on accuracy. Sparse architectures, such as mixture-of-experts (MoEs) [53, 36, 16] and sparse transformers [12, 6, 30], sparsify activations using sparsely-gated conditional computation, improving model capability without proportional increases in the memory or computational cost.

However, it is challenging for researchers and end users to take full advantage of sparsity in deep learning due to limitations in existing deep learning frameworks. Mainstream frameworks like PyTorch [48], JAX [9], and TensorFlow [1] provide sparse support for only a handful of operations on a few selected sparse data structures (tensor formats). As a result, leveraging sparsity in these frameworks often demands significant engineering effort to develop custom implementations for new operations. These frameworks lack the flexibility for researchers to explore the diverse range of potentially beneficial sparse data structures and computations in modern deep learning architectures.

Indeed, nearly every domain where sparsity arises has one-off libraries with hand-written and hand-optimized kernels, such as PyTorch Geometric [17] for graph learning, TensorLy [33] for sparse tensor factorization, and MegaBlocks [20] for MoEs. Although these libraries effectively target their specific applications, they do not generalize to other domains, resulting in a fragmented software ecosystem for sparse deep learning. The different abstractions create barriers to adoption, as they require users to learn new APIs and paradigms for each library. This also results in duplicated developer effort, as each domain-specific library implements its own sparse kernels and optimizations. Therefore, there is an opportunity for a unified framework that efficiently handles sparse operations and representations across various domains in deep learning.

To tackle these challenges, we introduce Scorch, the first library to integrate comprehensive and efficient sparse tensor computation capabilities into PyTorch. Scorch adds sparsity support to existing PyTorch operations, allowing researchers and practitioners to introduce sparsity into their models by simply declaring one or more tensors to be sparse. For instance, weight matrices can be made sparse, sparse gating can be used in architectures like MoEs for conditional computation, and sparse adjacency matrices can be efficiently processed in graph neural networks (GNNs). By preserving PyTorch’s interface, Scorch enables the natural and seamless expression of sparse models without requiring a separate ecosystem of sparse operations. Our initial work focuses on accelerating compute operations on CPUs, which enables sparse inference workloads and lays a foundation for general sparse computing in PyTorch. We leave GPU acceleration and auto-differentiation as future work.

The key insight behind Scorch is that it is possible to provide comprehensive sparsity support with minimal API changes to existing tensor frameworks. Sparsity is just a property of the values of a tensor and should not require a separate programming model. opt that everything doable with dense tensors should also be possible with sparse tensors. Scorch is designed to enable this functionality with a single line of code: import scorch as torch. Scorch enables general sparsity support by integrating state-of-the-art sparse code generation infrastructure [31, 13, 52] into a modern ML framework. However, prior work on sparse compilation lacks several technical components necessary to achieve such minimal API changes without sacrificing performance. We identify several key missing components and make the following technical contributions:

• A fast auto-scheduler for optimizing loop ordering, fusion, and the insertion of temporary data structures in the generated sparse kernels.

• A tiling algorithm that optimizes sparse compute kernels to improve cache locality.

• A format inference algorithm that determines the output format for each operation based on the nature of the operation and the input formats.

We evaluate Scorch on models across a range of domains, including graph neural networks, sparse autoencoders, and sparse transformers. We demonstrate speedups of 1.05–5.78$\times$ over PyTorch Sparse. This paper illustrates how to design a high-performance end-to-end sparse ML library. By enabling the natural expression and efficient execution of sparse models, Scorch can facilitate greater research exploration and adoption of sparsity.

2 Related Work

Deep learning frameworks. Several existing frameworks aim to support deep learning models for training and inference, but native sparse support in mainstream frameworks like PyTorch [48], JAX [9], and TensorFlow [1] is limited in generality, efficiency, and usability. Their sparsity support focuses on select formats like COO (coordinate list) and CSR (compressed sparse row), and a limited number of sparse operations like SpMM and SpMV. While libraries like MKL-Sparse and cuSPARSE provide low-level sparse primitives, they do not support high-order tensor operations, most formats across all operations, or fusion. Consequently, libraries that build support for general sparse tensor operation on top of them suffer from excessive data movement, which leads to poor performance.

Domain-specific libraries like PyTorch Geometric (PyG) [17] and Deep Graph Library (DGL) [57] have been developed to target sparse computation in specific applications. While they provide powerful capabilities within their target domains, they lack the generality to support other applications and sparse data structures. Scorch introduces a unified sparse tensor abstraction and compiler infrastructure to enable efficient and composable sparse deep learning across multiple domains.

Tensor algebra and machine learning compilation. Many compiler frameworks have been proposed for optimizing tensor operations and neural networks, such as Halide [49], TVM [11], XLA [22], Glow [50], and Tensor Comprehensions [55]. These compilers introduce new intermediate representations (IRs) and tensor algebra compilation techniques. However, they primarily focus on dense tensor computations and lack first-class support for sparse code generation.

Recent work on sparse tensor algebra compilers, such as TACO [31], Tiramisu [5], MLIR Sparse [7], and SparseTIR [60], has made significant progress in code generation for sparse tensor algebra expressions. However, these compilers are not integrated with deep learning frameworks, making it challenging to use them for sparse deep learning workloads. They also lack efficient and general automatic optimization algorithms necessary for high-performance sparse computation. Scorch incorporates ideas from the sparse code generation literature into the PyTorch ecosystem and adds automatic optimizations machinery to provide a seamless, end-to-end approach to productive and performant sparse deep learning.

3 Design

Scorch’s design is guided by three principles that aim to balance performance, usability, and generality in the context of sparse deep learning:

Seamless integration with PyTorch. To facilitate adoption and usability, a sparse deep learning library must integrate seamlessly with popular deep learning frameworks. Scorch is designed as a natural extension of PyTorch, with a familiar and intuitive interface for defining and manipulating sparse tensors. Users can leverage their existing PyTorch knowledge and codebases and easily switch between dense and sparse computation. This enables smooth adoption of Scorch and allows researchers to focus on develo** new models rather than learning new frameworks.

Unified abstractions for dense and sparse computation. Deep learning computations often involve a mix of dense and sparse tensors, but existing frameworks provide limited support for sparse tensors in their operators. For instance, PyTorch’s einsum and reshape operations do not support sparse tensors. To address this limitation, a sparse library should offer comprehensive operators that can freely operate on any mix of dense and sparse tensors, allowing users to express complex computations in natural tensor notation. This includes element-wise and non-linear operations, linear algebra, einsum operations, and tensor manipulations like reshape and transpose. Scorch accomplishes this by extending popular PyTorch operations to support sparse tensors. By providing unified operator support, Scorch lets users leverage the full power of PyTorch’s API while taking advantage of efficient sparse computation under the hood.

Robust automatic performance optimization. Achieving high performance on sparse workloads requires careful co-optimization of storage, loop orderings, use of intermediates, and tiling schemes. These choices collectively determine both the asymptotic complexity of a sparse kernel and the constant factors contributing to the runtime. The compiler stack of Scorch is capable of making reasonable decisions across this scheduling space using a set of carefully designed heuristics. Scorch’s optimization algorithms are designed around the following hierarchy of goals: First, format inference and loop ordering decisions should match those implemented in the state-of-the-art hand-written kernels. Second, loop ordering and storage decisions should ensure good worst-case asymptotic complexity to avoid unpredictable performance cliffs. Third, tiling should be applied to dense loops within sparse kernels to optimize cache usage. This heuristic hierarchy enables Scorch to match the performance of existing hand-written kernels and achieve predictable performance on unknown kernels, without burdening the user with manual performance tuning.

3.1 Programming Model

Scorch introduces a unified programming model for sparse and dense deep learning. Users opt into sparse execution via an import statement that make standard PyTorch operations compatible with sparse tensors. Figure 2 illustrates the simplicity and flexibility of Scorch’s programming model and how it improves productivity without sacrificing performance. The first example shows a Sampled Dense-Dense Matrix Multiplication (SDDMM) operation, commonly used in GNNs and recommender systems. It takes a sparse matrix $A$ and two dense matrices $B$ and $C$ as inputs and computes the sparse matrix $D$ as the element-wise product of $A$ and the matrix multiplication of $B$ and $C$. PyTorch does not support fused SDDMM for sparse COO inputs due to the lack of a hand-written kernel. Scorch enables users to express this computation concisely using einsum. By leveraging its compiler, Scorch automatically generates an efficient fused kernel for SDDMM. The generated kernel co-iterates over the sparse and dense tensors, avoiding the asymptotically more expensive dense matrix multiplication in the worst case. As seen in this example, fusion is crucial for performance, as it can reduce the overall computational complexity. By supporting mixed sparse and dense inputs operators like einsum, Scorch unlocks a wide range of sparse operations that would be difficult or impossible to express using PyTorch’s fixed set of hand-written sparse kernels.

Figure 1: Scorch SDDMM and sparse attention examples

Figure 2: Scorch architecture

3.2 System Overview

Figure 2 shows an overview of the architecture of Scorch. Its compiler lowers high-level sparse tensor operations into optimized kernels for efficient execution. Compilation begins by transforming Python code to a domain-specific intermediate representation (IR) that exposes optimization opportunities. The compiler then applies the following optimization passes to the IR:

• Automatic scheduling: Applies loop nest optimizations, tiling, and parallelization techniques specifically tuned for sparse tensor computations.

• Format inference: Automatically determines the most suitable output storage format for each operation based on its inputs and the operation.

• Code generation: Translates the optimized IR into efficient low-level kernels specialized for the specific storage formats of the input and output tensors.

A lightweight dispatcher orchestrates the execution of kernels. For common operations like SpMM, the dispatcher retrieves pre-compiled kernels from cache to save compilation time. For kernels that are not in the cache, the dispatcher invokes the above compiler at runtime to generate C++ code, which is then compiled and dynamically linked into the library.

4 Optimizations

Scorch introduces several optimizations to ensure the generated code has good and predictable performance. Efficient execution of sparse tensor operations requires careful consideration of the loop order, tiling, and the use of temporary tensor data structures. Suboptimal choices can lead to poor cache utilization, inefficient memory access patterns, and unnecessary data movement. In fact, the wrong loop order in a sparse loop nest can even change its asymptotic complexity [3]. To address this challenge, Scorch introduces new algorithms for auto-scheduling sparse kernels that determines the loop order, workspace insertion, and tiling strategy.

4.1 Automatic Loop Ordering and Workspace Insertion

Unlike dense computations, where different loop orders have the same asymptotic complexity, the choice of loop order can impact the worst-case asymptotic complexity and thus the observed performance of sparse computations [3]. Our loop ordering algorithm is designed to provide predictable performance to users by avoiding loop orders with poor asymptotic complexity. However, optimizing solely for the worst-case complexity may lead to bad average-case performance. We provide a fast heuristic auto-scheduler that achieves good average-case performance while avoiding the most asymptotically inefficient loop orderings.

Consider the sparse matrix-matrix multiplication (SpGEMM), $C_{ik}=\sum_{j}A_{ij}B_{jk}$, where the inputs $A$, $B$ and the output $C$ are stored in CSR format. Figure 4 illustrates three possible loop orderings for SpGEMM: outer product, Gustavson, and inner product. The loops iterate over the coordinates of different dimensions of the matrices, and the arrows denote the data structures for each matrix. For example, the $A$ matrix in Gustavson stores $i$ coordinates (denoted $A_{i}$) followed by $j$ coordinates (denoted $A_{j}$). The inner product algorithm has worse asymptotic complexity than the other two algorithms because the data structures of neither $A$ nor $B$ connects the coordinates in the two outer loops. They thus have to be iterated densely, meaning every component of the result has to be visited, even those that result in zeros. The other algorithms do not have this problem and thus perform much better. Moreover, algorithms that perform intersections between sparse matrices in higher loops (early filtering) perform better than those that perform intersections in inner loops. Scorch’s loop ordering algorithm avoids bad asymptotic behavior, even at the expense of transposes to allow matrices to be iterated in different orders.

Figure 3: SpGEMM loop orders

Figure 4: Tiling SpMM ($C_{ik}=\sum_{j}A_{ij}B_{jk}$)

Our algorithm (Algorithm 1) first collects and sorts index variables in the tensor expression by descending sparsity level. The sparsity level of an index variable is approximated by the presence of intersections (filters) between levels. In this case, $j\in A_{j}\cap B_{j}$ is the only sparse filter as it involves an intersection between the sparse level $A_{j}$ and the dense level $B_{j}$. The tie between $i$ and $k$ are broken by choosing the order that requires fewer transposes. This gives an initial loop order $\mathcal{L}=[j,i,k]$.

The algorithm then takes a greedy approach and iteratively evaluates the cost of pushing a sparse filter down the loop nest. In the SpGEMM example, the algorithm considers the net cost of moving $j$ to $\textit{pos}=1$: the filter no longer applies to the rows of $A$ in the $i$ loop, but the 2D workspace can be replaced by a 1D workspace since we no longer scatter into the $i$ dimension, and $A$ no longer needs to be transposed to iterate in $j,i$ order. As the net cost is negative, the loop order is updated to $[i,j,k]$. In the next iteration, we consider the net cost of moving $j$ to $\textit{pos}=2$, i.e., the inner product order $[i,k,j]$: while the 1D workspace can be eliminated since we no longer scatter into the result, the filter no longer applies to any of the tensors, and $B$ needs to be transposed to be iterated in the $k,j$ order. Moreover, as there are no input data structures that allow us to iterative sparsely from $i$ to $k$, it incurs a large penalty in the cost function. This makes net cost positive, so the loop order remains unchanged.

Algorithm 1 Loop Ordering and Workspace Insertion

1:Input: Tensor expression ${E}$, input tensors $\mathcal{T}$, output tensor $T_{out}$.

2:Output: Loop order $\mathcal{L}$ for efficient computation.

3:$\mathcal{I}\leftarrow\textsc{GetIndexVariables}(E)$ $\triangleright$ Set of all index variables in the expression

4:$\mathcal{S}\leftarrow\textsc{SortBySparsity}(\mathcal{I},E)$ $\triangleright$ Sort index variables by descending sparsity

5:$\mathcal{L}\leftarrow\textsc{InitLoopOrder}(\mathcal{S})$

6:for all $i\in\mathcal{S}$ do $\triangleright$ Loop ordering

7:     $\textit{currentPos}\leftarrow\textsc{GetPosition}(\mathcal{L},i)$ $\triangleright$ Current position of index $i$ in $\mathcal{L}$

8:     for pos $\leftarrow\textit{currentPos}$ to $|\mathcal{L}|+1$ do

9:         if Cost($\mathcal{L},i,\textit{pos})<0$ then $\triangleright$ Net cost of pushing $i$ to position pos

10:              $\mathcal{L}\leftarrow\textsc{MoveToPosition}(\mathcal{L},i,\textit{pos})$               

11:$G\leftarrow\textsc{InitGraph(}\mathcal{I}\textsc{)}$ $\triangleright$ Initialize a directed graph with nodes $\mathcal{I}$

12:for all $T\in\mathcal{T}\cup\{T_{out}\}$ do

13:     $\pi\leftarrow\textsc{GetModeOrdering(}T\textsc{)}$ $\triangleright$ Get the mode ordering of tensor $T$

14:     for $k\leftarrow 1$ to $|\mathcal{I}(T)|-1$ do

15:         $\textsc{AddEdgeToGraph(}G,\pi[k],\pi[k+1]\textsc{)}$ $\triangleright$ Add edge to capture ordering constraint      

16:while $\textsc{ContainsCycles(}G\textsc{)}$ do

17:     $\textsc{RemoveCheapestEdge}(G)$ $\triangleright$ Remove the cheapest edge from the graph

18:     $\mathcal{L}\leftarrow\textsc{UpdateLoopOrder(}\mathcal{L},e\textsc{)}$ $\triangleright$ Update the loop order by transposing the tensor

19:if $\textsc{HasSparseLevels(}T_{out}\textsc{)}$ then $\triangleright$ Workspace insertion

20:     $\mathcal{I}_{red}\leftarrow\textsc{GetReductionVariables(}\mathcal{T},T_{out}\textsc{)}$ $\triangleright$ Get reduction variables

21:     if $\textsc{ShouldInsertWorkspace(}\mathcal{I}_{red},\mathcal{I}(T_{out}),\mathcal{L}\textsc{)}$ then

22:         $W\leftarrow\textsc{InsertWorkspace(}i\textsc{)}$ $\triangleright$ Insert a workspace over loop $i$

23:         $(\mathcal{L}_{p},\mathcal{L}_{c})\leftarrow\textsc{SplitLoopOrder(}\mathcal{L},i\textsc{)}$ $\triangleright$ Split the loop order

24:         $\mathcal{L}_{p}\leftarrow\textsc{UpdateProducerLoop(}\mathcal{L}_{p},W,E\textsc{)}$ $\triangleright$ Update the producer loop

25:         $\mathcal{L}_{c}\leftarrow\textsc{UpdateConsumerLoop(}\mathcal{L}_{c},T_{out},W\textsc{)}$ $\triangleright$ Update the consumer loop      

26:return $\mathcal{L}$

4.2 Automatic Tiling

Tiling is a crucial optimization that improves cache utilization and reduces memory traffic by partitioning the iteration space into smaller blocks (tiles) that fit in the cache. Selecting which loops to tile is challenging, as suboptimal choices may hurt performance. Scorch introduces a novel sparse tiling algorithm that analyzes the tensor expression to determine which loops to tile based on several key observations:

First, opportunities for tiling can be inferred from the index variables present in each tensor access. In the SpMM example $C_{ik}=\sum_{j}A_{ij}B_{jk}$, there are three tensor accesses: $C[i,k]$, $A[i,j]$, and $B[j,k]$. If a tensor access is missing an index variable present in the full expression, it means that tensor access is reused across the loop corresponding to the missing index. For instance, $C[i,k]$ is missing $j$, indicating that it is reused across the $j$ loop. Thus, tiling the $i$ and $k$ loops would improve cache locality for $C$ if everything were dense. Second, tiling sparse dimensions is often counterproductive, as it requires performing expensive searches in the sparse data structures. Although it may be beneficial in highly tuned systems, Scorch avoids it in order to provide robust and predictable performance. In SpMM, for a CSR input $A$, the $j$ loop should not be tiled since $A_{j}$ is a sparse dimension. Third, not all dense dimensions should be tiled, as illustrated in Figure 4. While tiling $i$ in addition to $k$ would improve reuse of the rows of $A$, it would also require a larger working set for $B$ since the tiles of $k$ would be exhausted more frequently. In the dense case, the $j$ loop can be tiled to offset this, but since sparse dimensions are not tiled, tiling $i$ should be avoided. In general, tiling loops whose index variables are parent index variables of a sparse dimension should be avoided.

These insights form the basis of Scorch’s tiling algorithm, which can be summarized as follows: First, for each tensor access, collect its index variables if the set of index variables in the tensor access is a strict subset of the set of all index variables in the expression, as they indicate data reuse. Then, remove any index variables that correspond to a sparse dimension of any tensor. Finally, loops whose index variables are parents of a sparse index variable should not be tiled. The full algorithm is provided in Section C.1.

4.3 Format Inference

Tensor format inference is another key optimization in Scorch that automatically determines the storage format of the output tensor based on the storage formats of the input tensors and the semantics of the tensor expression. By inferring storage formats amenable to good performance, Scorch can generate efficient sparse kernels without requiring the user to manually specify the output format.

Scorch’s tensor format inference algorithm operates on a per-dimension basis, considering the interaction between the storage formats of each dimension in the input tensors. The intuition behind the algorithm is that multiplying a sparse level with any other level (sparse or dense) results in a sparse level in the output, while adding a dense level to any other level yields a dense level in the output. Appendix D contains details of the algorithm and its theoretical properties.

5 Evaluation

In this section, we demonstrate that Scorch adds general sparse support to PyTorch and achieves performance comparable to or better than hand-tuned sparse kernels on end-to-end models across several domains: graph neural networks, sparse autoencoders, and sparse transformers. We perform the benchmarks on an Apple M1 Ultra CPU (3.2 GHz, 20 cores) with 64 GB of memory.

5.1 Graph Neural Networks

We evaluate Scorch on the task of node classification using Graph Convolutional Networks (GCNs) [29]. GCNs are a popular class of graph neural networks that learn node representations by iteratively aggregating features from neighboring nodes. The core operation in a GCN layer is:

$$\mathbf{H}^{(l+1)}=\sigma(\hat{\mathbf{A}}\mathbf{H}^{(l)}\mathbf{W}^{(l)})$$ (1)

where $\mathbf{H}^{(l)}$ are the node features at layer $l$, $\hat{\mathbf{A}}$ is the normalized sparse adjacency matrix, $\mathbf{W}^{(l)}$ are the learnable weights, and $\sigma$ is a non-linear activation function.

We compare the inference performance of Scorch against implementations in PyTorch, PyG, and DGL on four benchmark datasets: Cora [43], Citeseer [21], PubMed [51], and OGBN-arXiv [26].

Figure 5 shows the inference times of the GCN model and the normalized speedups relative to PyTorch on the four datasets. Scorch achieves an average speedup of 2.08$\times$. The speedup achieved by Scorch can be attributed to several key optimizations in its SpMM kernel. Scorch applies tiling to improve cache locality and reduce memory traffic. It also uses pragma unroll directives to encourage loop unrolling, which can help reduce loop overhead and improve instruction-level parallelism.

Notably, the performance of PyG and DGL relative to PyTorch varies with the problem size. For smaller datasets like Citeseer and Cora, PyG and DGL achieves significant speedups over PyTorch. However, as the problem size grows, their performance advantage diminishes, for larger datasets like PubMed and OGBN-arXiv. This can be explained by the different algorithms used by these libraries for neighborhood aggregation. Instead of using SpMM directly, the gather-scatter algorithm indexes the source node features with the edge indices and scatters results into the destination nodes. While the gather-scatter algorithm is faster than SpMM for smaller graphs, it does not scale well to larger problems, as evident in the decreasing normalized speedups of PyG and DGL relative to PyTorch.

Figure 5: GCN Performance

5.2 Sparse Autoencoders

We evaluate the performance of sparse autoencoders implemented using Scorch and compare it against the PyTorch implementation. Autoencoders are unsupervised learning models that learn efficient data representations by encoding the input into a lower-dimensional latent space and then reconstructing the original input from the encoding. Sparse autoencoders introduce sparsity in the latent space, which can improve the interpretability and generalization of the learned representations [47]. The sparsity is typically achieved through regularization techniques like L1 penalty or KL-divergence loss. We train and test the autoencoders on four benchmark datasets: MNIST [35], CIFAR-10 [34], CIFAR-100 [34], and CelebA [39].

The sparse autoencoder architecture consists of an encoder with a sparse linear layer followed by ReLU activation, and a decoder with a dense linear layer followed by sigmoid activation. The models are trained using the mean squared error (MSE) loss and the Adam optimizer with a learning rate of 0.01. We train the models for 10 epochs with a batch size of 64.

Figure 6 shows the speedup achieved by Scorch over PyTorch for sparse autoencoder inference on the different datasets. Scorch achieves speedups of 1.42$\times$ to 5.78$\times$ over PyTorch, with the speedup being more pronounced on larger datasets like CIFAR-100 and CelebA. This is because the sparse linear layer in the encoder dominates the computation time, and Scorch’s optimized sparse kernels provide acceleration over PyTorch’s sparse linear algebra routines.

Figure 6: Sparse Autoencoder Performance

5.3 Sparse Transformers

Figure 7: Sparse Transformer Performance

We evaluate Scorch on the task of text classification using sparse transformer models, specifically the BigBird architecture [61]. BigBird leverages sparse attention mechanisms, including global, window, and random attention, to scale to long sequences with linear memory complexity. We benchmark the inference performance of BigBird implemented using Scorch versus the standard dense PyTorch implementation on three datasets: IMDB [42], AG News [64], and Yahoo Answers [64].

Figure 7 shows the inference times of BigBird. Scorch achieves an average speedup of 1.24$\times$, 1.14$\times$, and 1.05$\times$ in wall clock time for the three datasets respectively compared to PyTorch. Scorch’s ability to generate efficient kernels for sparse attention operations allows it to accelerate inference on real-world text classification tasks with varying sequence lengths and sparsity patterns.

5.4 Standard Kernels

We evaluate the performance of Scorch on several standard sparse kernels: Sparse Matrix-Vector Multiplication (SpMV), Sparse Matrix-Matrix Multiplication (SpMM), Sparse Matrix-Sparse Matrix Multiplication (SpMSpM), and Sampled Dense-Dense Matrix Multiplication (SDDMM). These operations are fundamental building blocks in many sparse deep learning models. We benchmark on matrices from the SuiteSparse Matrix Collection [14].

Figure 8 shows the absolute runtime of Scorch and torch.sparse for the four kernels across a range of matrix sizes and sparsity levels. For SpMV and SpMSpM, Scorch achieves performance similar to PyTorch as these operations have limited opportunities for optimization. However, for SpMM, Scorch exhibits better scaling and outperforms PyTorch for all but the smallest problems ($<10^{2}$ nonzeros). Scorch’s ability to generate efficient kernels becomes more advantageous as the problem size grows. More significant speedups are observed for SDDMM, where Scorch is orders of magnitude faster than PyTorch on bigger problems ($>10^{3}$ nonzeros). PyTorch cannot fuse the SDDMM operation and must perform a dense matrix multiplication followed by an element-wise multiplication with a sparse matrix. In contrast, Scorch is able to generate a single fused kernel that is asymptotically faster.

These results demonstrate the performance benefits of Scorch’s sparse code generation and optimizations, especially for operations like SDDMM where there are more opportunities for optimization. Scorch relieves users of the burden of handwriting and hand-optimizing sparse kernels while offering competitive performance. Appendix E contains additional details on the experimental setup.

Figure 8: Performance on key sparse operations

6 Limitations and Conclusion

Limitations. While Scorch supports a wide range of sparse tensor operations, it does not yet provide complete coverage of all possible operations. In particular, Scorch does not yet support automatic differentiation with sparse tensors, which is a problem outside the scope of this paper but a direction for future work to enable end-to-end training of sparse models. The current implementation of the compiler stack supports lowering to CPU-optimized C++ code. Additional performance may be achieved by lowering to Triton or CUDA code, and GPU code generation is left as future work.

Conclusion. Scorch is a PyTorch-based framework that accelerates sparse deep learning by providing efficient sparse computation capabilities and seamless integration with PyTorch. Scorch enables users to leverage the benefits of sparsity with minimal modifications to their existing code. Experimental evaluation demonstrates Scorch’s effectiveness in accelerating sparse ML workloads across various domains, achieving significant speedups over PyTorch Sparse. With its generality and seamless PyTorch integration, Scorch makes using sparsity in deep learning more accessible and paves the way for more exploration of efficient sparse models.

Appendix A Compilation

Scorch is designed as a layer on top of PyTorch that adds sparse functionality to existing operations. Scorch intercepts calls to PyTorch operations and dispatches them to either the traditional PyTorch infrastructure or to the Scorch compiler.

Scorch dispatches dense computations to the existing PyTorch computing infrastructure. This includes both tensor operations that operate on fully dense tensors and the fully dense sub-computations of blocked-sparse operations. By reusing the existing dense functionality, we make it easier for users to transition to Scorch, as dense operations continue to execute in the same execution environment and with the same performance as before. Moreover, we reduce the complexity of the Scorch compiler, as it does not need to be tuned to match existing dense computing libraries on completely dense operations.

For kernels with at least one sparse tensor, Scorch first performs format inference to decide the format of the output tensor. Then, Scorch generates a fused loop-level intermediate representation (IR) along the lines of TACO’s Concrete Index Notation IR [32]. The code generation algorithm for the loop-level IR also optimizes loop ordering, as described in Section 4.1, inserting intermediate temporary tensors as necessary. Unlike in dense tensor computations, the loop order of a sparse tensor computation affects its asymptotic complexity [3]. That is, changing a matrix multiplication from an inner product algorithm to a linear combination of row algorithm can lead to an arbitrarily large speedup. Loop ordering optimization is therefore critical and one of the contributions in Scorch is an algorithm that can quickly find a loop order that avoids a bad asymptotic complexity.

Next, Scorch tiles any dense loops inside sparse computations to improve cache utilization. Such loops appear because a sparse tensor computation often also contains dense tensors (e.g., SpMM and SDDMM) and because a sparse tensor may have some dimensions that are stored in a dense array (e.g., the first dimension of a CSR matrix is dense). The tiling algorithm for dense loops in Scorch (Section 4.2) is designed to avoid introducing costly random searches in sparse data structures. For example, Scorch does not tile sparse loops, although that can sometimes lead to performance improvements, because that would introduce costly searches in sparse data structures. It is difficult for a compiler to determine whether the cache benefits outweighs the searching cost. We therefore designed Scorch to not tile sparse loops so that the compiler produces predictable good performance.

Finally, the tiled loops are lowered to C++ code, JIT compiled and linked via PyTorch’s custom C++ extensions loader. A high-level description of this process is provided in Algorithm 2.

Algorithm 2 Scorch Kernel Compilation

1:Input: Tensor expression ${E}$ with input tensors $\mathcal{T}$ and untyped output tensor $T_{out}$.

2:if all $\mathcal{T}$ are dense then

3:     return torch.compile(${E}$) $\triangleright$ Dynamic dispatch for dense kernels.

4:$T_{out}$.format $\leftarrow$ InferFormat(${E}$, $\mathcal{T}$) $\triangleright$ Appendix D

5:$\mathcal{L}\leftarrow$ CompileToLoops(${E}$, $\mathcal{T}$, $T_{out}$) $\triangleright$ Algorithm 1

6:$\mathcal{L}_{\mathrm{tiled}}\leftarrow$ TileOperations(${E}$, $\mathcal{T}$, $\mathcal{L}$) $\triangleright$ Algorithm 3

7:return LowerToC($\mathcal{L}_{\mathrm{tiled}}$)

The current implementation of Scorch compiles a single user-provided expression at a time and does not fuse across expressions. We leave extending our compiler infrastructure to fuse sparse operations across expressions. Such work could, for example, use a tracing facility such as the one introduced in PyTorch 2.0 [4] to capture the computational graph that serves as the basis for global optimization.

Appendix B Implementation Details

Scorch is implemented as a Python library, with its compiler written in Python as well. The Scorch compiler generates C++ code and pybind bindings, which are then compiled and loaded using PyTorch’s custom C++ extension machinery.

B.1 Sparse Tensor Representation

Scorch represents sparse tensors using a unified abstraction that can represent different underlying physical sparse storage formats like COO, CSR, and DCSR. This allows specifying computations independently of storage details, leaving the task of generating low-level code that interacts with specific data structures to the compiler.

The Scorch implementation has two primary classes that manage sparse tensors. A tensor class contains general information about the tensor that is independent of specific storage, such as its name, shape, and component type. The tensor class also has a reference to an object of a storage class. A tensor storage contains the physical storage of a tensor, including a storage descriptor (i.e., a format), arrays storing sparse indices that identify the stored (typically non-zero) values of a tensor, such as coordinate lists and CSR row pointers, and an array containing the stored values of the tensor.

B.2 Sparse Tensor Operations

Scorch supports arbitrary tensor contractions and element-wise operations on sparse tensors. This includes, but is not limited to:

• sparse matrix-vector multiplication,

• sparse matrix-matrix multiplication,

• tensor contractions,

• element-wise sparse addition, subtraction, multiplication, and

• fused operations.

At the core, Scorch can compile generalized einsum contraction operations. The data structures that store each tensor in these operations can be determined separately. The code generator builds on prior work on sparse tensor algebra code generation from the TACO line of work [31, 32]. Such code generators generate code that co-iterates over any number of data structures stored in different formats. Thus, they support generating fused code that operates on tensors stored in disparate formats.

B.3 Sparse Temporary Workspaces

One key feature in Scorch that enables general sparse computations is the support for sparse temporary tensors, often called sparse workspaces [62]. Such temporaries are useful in kernels where the loop order causes scattering into the result tensor, but where the result tensor data structure does not support random inserts (e.g., a linear combination of rows matrix multiplication with a CSR result). The sparse temporary tensors in Scorch allow computations to scatter and gather results in arbitrary order into a sparse output tensor of any format, while also offering significant memory efficiency by storing temporary intermediate tensors in a compressed format, avoiding the need to materialize intermediate zeros values as in dense tensors. In Scorch, a red-black tree is used to store intermediate non-zero values. We chose red-black trees because they are simple and inherently provide the features we need, namely the ability to insert coordinates-value pairs and the ability to iterate over them in sorted order. Workspaces improve performance by avoiding unnecessary format materialization and compression.

Appendix C Auto-scheduling Algorithms

C.1 Tiling

The tiling algorithm analyzes the tensor expression and determines which loops to tile based on several key observations, as described in Section 4.2. The pseudocode for the tiling algorithm is provided in Algorithm 3.

Algorithm 3 Tiling Sparse Tensor Operations

1:Input: A tensor expression $E$ containing input tensors $\mathcal{T}$, and a loop structure $\mathcal{L}$.

2:Output: Tiled loop structure $\mathcal{L}_{\mathrm{tiled}}$ for efficient computation.

3:$\mathcal{W}\leftarrow\emptyset$ $\triangleright$ Initialize the working set of index variables

4:$\mathcal{S}\leftarrow\textsc{GetIndexVariables(}E\textsc{)}$ $\triangleright$ Get the set of all index variables in the expression

5:for each tensor access $T[\mathcal{I}]$ in the expression $E$ do

6:     if $\mathcal{I}\subset\mathcal{S}$ then $\triangleright$ Check if indices of $T$ are a strict subset of $\mathcal{S}$

7:         $\mathcal{W}\leftarrow\mathcal{W}\cup\mathcal{I}$ $\triangleright$ Add indices to the working set      

8:for all $i\in\mathcal{W}$ do $\triangleright$ Remove sparse index variables

9:     if there exists a tensor $T\in\mathcal{T}$ where the dimension of $T$ corresponding to $i$ is sparse then

10:         $\mathcal{W}\leftarrow\mathcal{W}\setminus\{i\}$      

11:for all $i\in\mathcal{W}$ do $\triangleright$ Remove index variables that are parents of sparse dimensions

12:     if there exists a sparse index variable $j$ such that $i$ is a parent of $j$ in the loop nest $\mathcal{L}$ then

13:         $\mathcal{W}\leftarrow\mathcal{W}\setminus\{i\}$      

14:for all $i\in\mathcal{W}$ do $\triangleright$ Tiling

15:     $(i_{\text{outer}},i_{\text{inner}})\leftarrow\textsc{TileIndex}(i)$ $\triangleright$ Split loop $i$ into outer and inner loops

16:     $\mathcal{L}\leftarrow\textsc{ReorderLoops}(\mathcal{L},i_{\text{outer}},0)$ $\triangleright$ Reorder $i_{\text{outer}}$ to outermost in the loop nest

17:return $\mathcal{L}_{\mathrm{tiled}}\leftarrow\mathcal{L}$

Appendix D Tensor Formats and Format Inference

Scorch uses the sparse tensor format description language proposed by Kjolstad et al. [31] and Chou et al. [13]. When users create a PyTorch tensor, it is dense by default as it is in vanilla PyTorch. Scorch extends the PyTorch tensor constructor to allow the user to specify that a tensor should be sparse. Sparse tensors are stored in the coordinate format by default, but also lets users specify a specific tensor formats. The Scorch tensor format description language is flexible and support many different types of data structures, such as compressed vectors, compressed matrices, compressed tensors, blocked-sparse matrices and tensors, dense tensors, and coordinate matrices and tensors.

Although users must specify whether a tensor defined using a constructor is sparse, Scorch automatically infers the (dense or sparse) data structures of tensors that result from some operations. So if the user multiplies two sparse matrices (SpGEMM), Scorch will infer that the result should be stored in a sparse matrix, as well as the specific sparse data structure. And if the user multiplies a sparse matrix with a dense matrix (SpMM), Scorch automatically infers that the result should be stored in a dense array.

D.1 Tensor Formats

A tensor format describes the data structures that store the non-zero elements of a tensor. Following Kjolstad et al. [31], we allow a separate description of the data structure of each dimension of a tensor. So a compressed sparse row (CSR) matrix stores the set of rows in a dense data structure and each row in a compressed data structure, while a doubly-compressed sparse row (DCSR) matrix stores both dimensions in a compressed data structure. We support three types of per-dimension data structures that can be composed any way to store a tensor of any dimensionality:

• dense: All elements, including zeros, are explicitly stored.

• compressed: Non-zeros are stored using compressed index arrays storing their coordinates.

• coordinate: Non-zeros are stored as a list of coordinate-value pairs.

The full format of a tensor is determined by both the ordering of dimensions and the data structure for each dimension. For example, a matrix (2D tensor) can be stored in many ways: a row-major dense matrix (dense rows, dense columns), a row-major compressed sparse row (CSR) matrix (dense rows, compressed columns), a column-major compressed sparse column (CSC) matrix (dense columns, compressed rows), a doubly compressed sparse row (DCSR) matrix (compressed rows, compressed columns), and a row-major coordinate list (COO) matrix (coordinate rows, coordinate columns).

Our way to specify coordinate tensors simplifies their specification in prior work by Chou et al. [13]. They introduced a singleton data structure that needed to be composed with a duplicate version of a compressed data structure, so a coordinate matrix became (unordered compressed rows, singleton columns). Scorch, on the other hand, represents coordinate tensors more uniformly using a coordinate data structure, so a coordinate matrix becomes (coordinate rows, coordinate columns).

D.2 Format Inference Algorithm

Given a tensor computation expressed as an einsum operation or as a general tensor algebra expression, we provide an algorithm to infer the format of the output tensor based on the formats of the input tensors. Our inference algorithm determines the storage format for each level of the output tensor independently and thus naturally extends to tensors of any dimensionality.

The inference algorithm traverses the tensor expression and predicts the format of the result by separately predicting the data structure of each of its dimensions. This is done using basic algebraic reasoning: the element-wise multiplication of a sparse vector by any other vector type should be sparse, since the result will be at least as sparse as either operand. And the element-wise addition of a dense vector with any other vector type should be dense, since the result will be at least as dense as either operand. Furthermore, when two sparse vectors are added, Scorch makes the result sparse. Although the result is the union of the operands, and thus denser than either operand, we elect to make Scorch conservative. The reason for such conservativeness is that instantiating a dense tensor when the result is very sparse can asymptotically increase memory and compute usage, while instantiating a sparse tensor when the result is sparse only increases memory and compute usage by a constant factor. Finally, when two dense vectors are multiplied, Scorch makes the result dense. We provide the inference algorithm in Algorithm 4, where compressed and coordinate formats are both considered Sparse.

Algorithm 4 Tensor Format Inference

1:procedure Infer($E$, $lvl$)

2:     match E

3:     | $A+B$ $\rightarrow$ match Infer($A$, $lvl$), Infer($B$, $lvl$) with

4:            | Dense, _ $\rightarrow$ return Dense

5:            | _, Dense $\rightarrow$ return Dense

6:            | _, _ $\rightarrow$ return Sparse

7:     | $A*B$ $\rightarrow$  match Infer($A$, $lvl$), Infer($B$, $lvl$) with

8:            | Sparse, _ $\rightarrow$ return Sparse

9:            | _, Sparse $\rightarrow$ return Sparse

10:            | _, _ $\rightarrow$ return Dense

11:     | $T$ $\rightarrow$ TensorFormat($T$, $lvl$)

12:     end

D.3 Example

Consider the following tensor computation:

$$D_{ij}=\sum_{k}A_{ik}B_{kj}+C_{ij}$$ (2)

Suppose the input tensors have the following formats:

• $A$: CSR (dense $i$, compressed $k$)

• $B$: CSR (dense $k$, compressed $j$)

• $C$: DCSR (compressed $i$, compressed $j$)

To infer the format of the output tensor $D$, we apply the format inference algorithm:

1. First, we consider the multiplication sub-expression $T_{ij}=\sum_{k}A_{ik}B_{kj}$:

   • For $T[i]$: dense $A[i]$ contributes through multiplication, so $T[i]$ should be dense.

   • For $T[j]$: compressed $B[j]$ contributes through multiplication, so $T[j]$ should be compressed.

2. Next, we consider the addition of $T$ and $C$ to form the final output $D$:

   • For $D[i]$: dense $T[i]$ contributes through addition, while compressed $C[i]$ contributes through addition. Since $T[i]$ is dense, $D[i]$ should be dense.

   • For $D[j]$: compressed $T[j]$ contributes through addition, while compressed $C[j]$ contributes through addition. Therefore, $D[j]$ should be compressed.

   Therefore, the final output tensor $D$ should be stored in CSR format.

This example demonstrates how the format inference algorithm handles expressions with a combination of multiplications and additions by recursively applying the inference rules and combining the results.

D.4 Discussion

The tensor format inference algorithm is implemented in the Scorch framework as part of the compiler pipeline. When a tensor computation is encountered, Scorch analyzes the formats of the input tensors and applies the format inference algorithm to determine an appropriate format for the output tensor(s).

The inferred formats are then used to guide the code generation process, ensuring that the appropriate sparse or dense kernels are generated. This automatic format inference capability lets Scorch optimize tensor computations based on the data structures of the input data, resulting in improved memory efficiency and computational performance.

By integrating format inference into the compiler workflow, Scorch enables the seamless mixing of sparse and dense tensors in deep learning models, without requiring manual specification of the result of each and every compute operation. This simplifies the development process and lets users focus on the high-level logic of their models while leveraging the benefits of sparse computations where appropriate.

One limitation of the current algorithm is that it does not consider potential trade-offs between storage efficiency and computational efficiency. In some cases, it may be beneficial to use a denser format for the output tensor to enable more efficient computation, even if it requires more storage. Extending the format inference algorithm to take into account these trade-offs is an interesting direction for future work.

Overall, format inference is a key component of Scorch that enables efficient and transparent sparse tensor computation. By automatically inferring the optimal sparse format, Scorch simplifies the user experience and ensures high performance across a wide range of sparse tensor operations and sparsity patterns.

Appendix E Experimental Details

All experiments were performed on an Apple M1 Ultra CPU (3.2 GHz, 20 cores) with 64 GB of memory.

E.1 Standard Kernels Benchmark Details

We run each benchmark 10 times and report the average runtimes in Figure 8. For SpMV, we perform a matrix-vector multiplication of each SuiteSparse matrix with a randomly generated dense vector. For SpMM, we perform the matrix multiplication of each SuiteSparse matrix with a randomly generated dense matrix. For SpMSpM, we truncate any non-square SuiteSparse matrices to be square and matrix multiply them with their transpose. For SDDMM, we perform the element-wise multiplication of each SuiteSparse matrix with the matrix multiplication of two randomly generated dense matrices.

Matrix Formats and Characteristics.

For the SpMV, SpMM, and SDDMM benchmarks, we use sparse matrices in the Compressed Sparse Row (CSR) format, which is widely supported across sparse libraries. However, for SpMSpM, PyTorch does not support matrix multiplication between two CSR matrices on non-Intel machines due to the absence of the Intel Math Kernel Library (MKL). The specific error message is: “addmm: computation on CPU is not implemented for SparseCsr + SparseCsr @ SparseCsr without MKL.”

To work around this limitation and still evaluate SpMSpM performance, we use the COO (Coordinate) format for SpMSpM. Additionally, since not all matrices in the SuiteSparse collection are square, we truncate non-square matrices to be square and multiply them with their transpose to obtain a valid SpMSpM operation.

PyTorch Performance on Small SpMM Problems.

We observe that PyTorch Sparse outperforms Scorch on SpMM for small problems with fewer than $10^{2}$ nonzeros. This is because PyTorch’s SpMM implementation is optimized for small matrices. Their gather-scatter approach is more efficient for small inputs but does not scale well to larger matrices. Scorch’s generated kernels have some overhead that dominates for small problems, but they scale much better to larger matrices, resulting in the speedups seen in Figure 8.

E.2 Graph Neural Networks

We evaluate the inference performance of Graph Convolutional Networks (GCNs) [29] on four node classification datasets: Cora [43], Citeseer [21], PubMed [51], and OGBN-arXiv [26]. The model architecture and hyperparameters are listed in Table 1. We use the Adam optimizer with a learning rate of 0.01 and weight decay of 5e-4 to train the model for 200 epochs with a batch size of 1 (full-batch). Dropout with a rate of 0.5 is applied to the hidden layer during training.

For each dataset, we train the models on the training set and evaluate the inference time and accuracy on the test set using an Apple M1 Ultra CPU with 64 GB of memory. The experiments are run in PyTorch 2.2.1, PyTorch Geometric 2.5.0, and DGL 2.0.0, and Scorch 0.1.0. We run each inference experiment 50 times with 5 warm-up runs and report the average speedups of various frameworks relative to PyTorch, as well as the absolute inference times, in Figure 5.

To ensure a fair comparison, we use the same GCN architecture across all frameworks. The PyTorch and Scorch implementations use a custom GCN layer, while PyG and DGL use their built-in GCN layers. The same model is trained with PyG and the train weights are loaded into the PyTorch, Scorch, PyG, and DGL models after adjusting for any differences in the parameter shapes and orderings.

Hyperparameter	Value
Hidden channels	128
Activation function	ReLU
Dropout rate	0.5
Optimizer	Adam
Learning rate	0.01
Weight decay	5e-4
Batch size	1 (full-batch)
Training epochs	200

Table 1: GCN model architecture and hyperparameters for node classification.

E.3 Sparse Autoencoders

We evaluate the performance of sparse autoencoders on four datasets: MNIST [35], CIFAR-10 [34], CIFAR-100 [34], and CelebA [39]. The datasets are preprocessed as follows:

• MNIST: The images are converted to tensors and flattened into a 1D vector of size 784 (28$\times$28).

• CIFAR-10 and CIFAR-100: The images are converted to grayscale, then to tensors, and flattened into a 1D vector of size 1024 (32$\times$32).

• CelebA: The images are converted to grayscale, resized to 64$\times$64, then to tensors, and flattened into a 1D vector of size 4096 (64$\times$64).

The sparse autoencoder architecture and hyperparameters are summarized in Table 2. The encoder is a sparse linear layer followed by ReLU activation, and the decoder is a dense linear layer followed by sigmoid activation.

Component	Details
Encoder	Sparse Linear (input size $\rightarrow$ 256)
Encoder Activation	ReLU
Decoder	Dense Linear (256 $\rightarrow$ input size)
Decoder Activation	Sigmoid
Loss Function	Mean Squared Error (MSE)
Optimizer	Adam
Learning Rate	0.01
Batch Size	64
Training Epochs	10

Table 2: Sparse autoencoder architecture and hyperparameters.

The models are trained using the mean squared error (MSE) loss and the Adam optimizer [28] with a learning rate of 0.01. We use a batch size of 64 and train the models for 10 epochs. The experiments are conducted on an Apple M1 Ultra CPU with 64GB of memory.

During inference, we measure the average reconstruction loss and the wall-clock time on the test set. The reconstruction loss is computed using the MSE loss between the input images and the reconstructed images. The wall-clock time includes the time taken to convert the input data to a sparse CSR tensor, forward pass through the model, and compute the loss.

The experiments are implemented in PyTorch 2.2.1 and Scorch 0.1.0. To ensure a fair comparison, we use the same model architecture, hyperparameters, and random seed across both frameworks. We run each inference experiment 50 times with 5 warm-up runs and report the average speedups of Scorch relative to PyTorch, as well as the absolute inference times, in Figure 6.

E.4 Sparse Transformers

Sparse transformers are variants of the Transformer architecture [56] that leverages sparse attention mechanisms to improve computational efficiency and scalability. In this evaluation, we implement BigBird [61], a specific type of sparse transformer that employs a combination of global, sliding window, and random sparse attention patterns.

BigBird is designed to handle long sequences while maintaining a manageable computational complexity. It achieves this by using a sparse attention mechanism that attends to a subset of tokens in the sequence, rather than attending to all tokens as in the standard transformer. The sparse attention pattern in BigBird consists of three components:

1. Global attention: Attend to all tokens in fixed-size blocks at regular intervals.

2. Sliding window attention: Attend to neighboring tokens within a fixed-size window that slides over the sequence.

3. Random attention: Attend to a fixed number of randomly selected tokens in the sequence.

By combining these attention patterns, sparse transformers can capture both local and global dependencies while significantly reducing the computational cost compared to the standard Transformer.

Sparse transformers have been successfully applied to various natural language processing tasks, such as text classification, question answering, and language modeling, where the input sequences can be very long. It has also shown promising results in other domains, such as genomics and time series analysis, where the ability to handle long sequences is crucial.

We evaluate the inference performance of the BigBird model [61] on three text classification datasets: AG News [64], IMDB [42], and Yahoo Answers [64]. The model architecture and hyperparameters are listed in Table 3. We use the AdamW optimizer [40] with a learning rate of 0.001 to train the models for 5 epochs with a batch size of 64. The sparse attention is configured to use a block size of 16, 2 global tokens, 2 random blocks, and 2 sliding blocks.

For each dataset, we train the models on the training set and evaluate the inference time on the test set using an Apple M1 Ultra CPU. The experiments are implemented in PyTorch 2.2.1 and Scorch 0.1.0. We run each inference experiment 10 times with 5 warm-up runs and report the average speedups of Scorch relative to PyTorch, as well as the absolute inference times, in Figure 7.

Hyperparameter	Value
Embedding size	128
Hidden size	128
Intermediate size	256
Number of hidden layers	2
Number of attention heads	4
Hidden activation	gelu
Attention dropout	0.1
Hidden dropout	0.1
Sparse attention block size	16
Sparse attention global tokens	2
Sparse attention random blocks	2
Sparse attention sliding blocks	2
Optimizer	AdamW
Learning rate	0.001
Batch size	64
Training epochs	5

Table 3: BigBird model architecture and hyperparameters for text classification.

Appendix F Extended Related Work

Sparse Tensor Formats. To avoid unnecessary computation and storage of zero values, various compressed sparse tensor formats have been developed. These aim to store only the meaningful non-zero values and indices required to fully represent a sparse tensor. Popular formats that store irregular tensors include compressed sparse row/column storage (CSR/CSC) which compactly stores non-zero values and their row/column indices. Coordinate format (COO) uses tuples of indices and corresponding values. Block sparse formats partition tensors into dense sub-blocks for efficiency. Compressed sparse fiber exploits skewed sparsity patterns by orienting storage along one dimension. These formats are all supported by Scorch.

There are also many tensor formats that are not supported by Scorch. Many of these have been developed to take advantage of additional structure in the tensors, such as ELLPACK, which assumes a bounded number of non-zeros per row, and the diagonal matrix format, which assumes all non-zeros are located on a small number of diagonals. Other formats store additional data to matrices in both row- and column-major order, such as the CSB matrix format [10] and the blocked matrix format used by MegaBlocks [20]. Finally, several tensor formats have been proposed in recent years to further improve the performance of tensor algebra operations at the expense of more complexity in the data structures, such as CSR5 [38] and HiCOO [37]. In the current version of Scorch, we chose to focus on the data structures that account for the vast majority of use, to manage engineering complexity. However, we believe it is interesting future work to add more formats into Scorch while supporting any operations on any combination of data structures..

Each format makes different trade-offs between compression, indexing overhead, and read/write efficiency for varied sparsity characteristics. Choosing appropriate sparse storage and operations can significantly reduce computational complexity and memory usage compared to dense defaults. However, current deep learning frameworks lack native support for sparse data structures across all compute operations.

Sparse Kernels. In addition to compressed storage, many specialized sparse kernels have also been created, including for sparse matrix multiplication (SpMM), sparse matrix-vector multiply (SpMV), sparse convolution, and more. Key optimizations in these kernels include techniques like iteration space tiling, loop reordering to improve locality, vectorization, and load balancing to handle irregular sparsity. Highly optimized sparse libraries like NVIDIA cuSPARSE and Intel MKL provide some essential sparse operations, but composing and scheduling sparse kernels in deep learning frameworks remains a challenging problem.

Despite the potential benefits of sparsity, adoption in deep learning still faces multiple key challenges: the lack of versatile software frameworks to natively exploit sparsity in the data and models; the mismatch between existing dense architectures like GPUs/TPUs and sparse computation; the difficulty of fusing sparse operations with other layers in an end-to-end model; the irregular memory access patterns of sparse kernels impairing performance compared to dense routines; the increased difficulty of parallelization and load balancing due to sparse connectivity; and finding the right abstraction level to balance performance and productivity.

Auto-scheduling. Scorch contributes techniques for automatically determining optimal schedules for sparse kernels. These are related to auto-scheduling approaches in compilers. The automatic scheduling of computational kernels is challenging due to the exponential search space of possible optimizations like loop transforms and data layouts. Prior auto-scheduling techniques for dense [2, 66, 65] and sparse tensors [3, 27] rely on cost modeling, optimization heuristics, and black-box search. Many auto-schedulers for dense tensor programs exist, such as the Halide auto-scheduler [2], FlexTensor [66], and Ansor [65]. Some recent work has focused on auto-scheduling for tensor algebra. CIN-P [3] enumerates schedules based on asymptotic costs to find optimal schedules for general sparse tensor algebra expressions. SpTTN-Cyclops [27] tunes contraction paths and index orders for SpTTN (contractions of a single large sparse tensor with several dense tensors) kernels using both enumeration and efficient search. Scorch contributes a lightweight heuristic-based auto-scheduler for general sparse kernels that leverages the structure of sparse tensor operations to efficiently explore the search space of different tensor kernels. Moreover, we show in our evaluation that the Scorch heuristic auto-scheduler is able to find good implementations of many operations.

Sparsity in Deep Learning. Sparsity is inherent in many emerging domains for deep learning, including recommender systems [46], drug discovery [63, 15], and web-scale graph analytics [23]. For example, in recommendation systems, user-item interaction matrices are often sparse, as each user interacts with a small subset of the catalog. Molecular graph representations are similarly sparsely connected. Large-scale graphs such as web and social network graphs also exhibit power law distributions, resulting in highly sparse adjacency matrices.

Despite this sparsity, the default approach in deep learning relies heavily on dense matrix multiplication and convolution operations [54]. Using dense matrix multiplication algorithms on extremely sparse data wastes both computation and memory. The massive computational and memory burdens incurred by ignoring sparsity quickly become prohibitive as data continues growing in scale and dimensionality [25].

Several works have developed sparse neural network architectures. Methods like sparse evolutionary training [44] and dynamic sparse reparameterization [45] induce sparsity during training through pruning or regularization. The lottery ticket hypothesis [18] shows that networks can be trained from scratch to be sparse. Structured sparsity techniques induce sparsity in regular patterns, such as block sparsity and channel sparsity, to optimize memory access and computation for specific hardware architectures [58]. While these methods demonstrate the potential benefits of sparsity in deep learning, they often require custom sparse kernels to be implemented for the specific sparse operations used [19]. Scorch enables researchers to more easily explore novel sparse architectures by providing a general framework for efficient sparse computation. Rather than needing to implement custom kernels, users can rely on Scorch’s compiler to automatically generate performant code for their particular sparse operations. This reduces the engineering burden and allows research to focus on modeling innovations.