Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Changing Base Without Losing Pace: A GPU-Efficient Alternative to MatMul in DNNs

Nir Ailon, Akhiad Bercovich, Yahel Uffenheimer, Omri Weinstein

TL;DR

The paper tackles the bottleneck of MatMul in deep networks by introducing Strassen-Tile (STL), a GPU-native bilinear operator that operates on tiles of weight and activation matrices through learnable encoders and a decoder. STL reduces FLOPs substantially while preserving or increasing parameter count, demonstrated via tile-based approximations of 4×4 MatMuls and improvements in Imagenet-1K accuracy for a SoTA ViT model, with observed wall-clock speedups on GPUs. The work grounds STL in Strassen normal forms, analyzes its FLOPs/IO complexity, and details a GPU-friendly implementation that decomposes the compute into per-tile MatMuls over encoded tiles. Empirically, STL shows promise in under-parameterized regimes, with initialization and training strategy playing crucial roles, and points toward future exploration on larger architectures and specialized kernels. Overall, STL is a compelling building block toward scalable and cost-efficient AI, balancing speed, accuracy, and parameterization.

Abstract

Modern AI relies on huge matrix multiplications (MatMuls), whose computation poses a scalability problem for inference and training. We propose an alternative, GPU native bilinear operator to MatMuls in neural networks, which offers a three-way tradeoff between: speed, accuracy and parameter count. In particular, this operator requires substantially fewer FLOPs to evaluate ($\ll n^3$), yet increases the parameter count compared to MatMul ($\gg n^2$). We call this operator Strassen-Tile (STL). The key idea behind STL is a local learnable change-of-basis, applied on tiles of the weight and activation matrices, followed by an element-wise product between the tiles, implemented simultaneously via MatMul. The key technical question we study is how to optimize the change-of-basis of a given layer, which is a highly non-convex problem. We show that theory-backed initializations (inspired by fast matrix and polynomial multiplication) lead to substantially better accuracy than random SGD initialization. This phenomenon motivates further algorithmic study of STL optimization in DNNs. Our experiments demonstrate that STL can approximate 4x4 MatMul of tiles while reducing FLOPs by a factor of 2.66, and can improve Imagenet-1K accuracy of SoTA T2T-ViT-7 (4.3M parameters) while lowering FLOPs. Even with non-CUDA optimized PyTorch code, STL achieves wall-clock speedups in the compute-bound regime. These results, together with its theoretical grounds, suggest STL as a promising building block for scalable and cost-efficient AI.

Changing Base Without Losing Pace: A GPU-Efficient Alternative to MatMul in DNNs

TL;DR

The paper tackles the bottleneck of MatMul in deep networks by introducing Strassen-Tile (STL), a GPU-native bilinear operator that operates on tiles of weight and activation matrices through learnable encoders and a decoder. STL reduces FLOPs substantially while preserving or increasing parameter count, demonstrated via tile-based approximations of 4×4 MatMuls and improvements in Imagenet-1K accuracy for a SoTA ViT model, with observed wall-clock speedups on GPUs. The work grounds STL in Strassen normal forms, analyzes its FLOPs/IO complexity, and details a GPU-friendly implementation that decomposes the compute into per-tile MatMuls over encoded tiles. Empirically, STL shows promise in under-parameterized regimes, with initialization and training strategy playing crucial roles, and points toward future exploration on larger architectures and specialized kernels. Overall, STL is a compelling building block toward scalable and cost-efficient AI, balancing speed, accuracy, and parameterization.

Abstract

Modern AI relies on huge matrix multiplications (MatMuls), whose computation poses a scalability problem for inference and training. We propose an alternative, GPU native bilinear operator to MatMuls in neural networks, which offers a three-way tradeoff between: speed, accuracy and parameter count. In particular, this operator requires substantially fewer FLOPs to evaluate (), yet increases the parameter count compared to MatMul (). We call this operator Strassen-Tile (STL). The key idea behind STL is a local learnable change-of-basis, applied on tiles of the weight and activation matrices, followed by an element-wise product between the tiles, implemented simultaneously via MatMul. The key technical question we study is how to optimize the change-of-basis of a given layer, which is a highly non-convex problem. We show that theory-backed initializations (inspired by fast matrix and polynomial multiplication) lead to substantially better accuracy than random SGD initialization. This phenomenon motivates further algorithmic study of STL optimization in DNNs. Our experiments demonstrate that STL can approximate 4x4 MatMul of tiles while reducing FLOPs by a factor of 2.66, and can improve Imagenet-1K accuracy of SoTA T2T-ViT-7 (4.3M parameters) while lowering FLOPs. Even with non-CUDA optimized PyTorch code, STL achieves wall-clock speedups in the compute-bound regime. These results, together with its theoretical grounds, suggest STL as a promising building block for scalable and cost-efficient AI.

Paper Structure

This paper contains 25 sections, 2 theorems, 17 equations, 5 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work.
  3. Structure of the Paper.
  4. Strassen Normal Forms
  5. Strassen-Tile Operator $\mathsf{STL}$
  6. FLOPs Complexity Analysis of $\mathsf{STL}$
  7. GPU-Friendly Implementation of the Element-Wise Product
  8. GPU Complexity Analysis
  9. Conclusion
  10. Experiments
  11. Class 0 Experiments - Synthetic $4\times 4$ matrices
  12. Class 1 Experiments - Training From Scratch with $\mathsf{STL}$
  13. Conclusion
  14. Parameter Increase in $\mathsf{STL}$
  15. Theoretical Foundations and Initialization
  16. ...and 10 more sections

Key Result

Corollary 3.3

Assuming $n\gg t$, the FLOP count of $\mathsf{STL}$ for square matrices is $O(rn^3/t^3)=O(n^3c/t)$.

Figures (5)

  • Figure 1: Observed speedup factor for $\mathsf{STL}$ for tile size $t=4$, $r=16,20,\ldots,48,49$, and matrices of size $n \times n$ with $n=4096, 8192, 16384$, using a (non-optimized) PyTorch implementation. As expected, the throughput speedup is almost linear in the tensor rank. (Note: $r=49$ can imitate exact matrix multiplication by Strassen, and is given here for completeness.)
  • Figure 2: The blue line is estimation of the error of $\mathsf{STL}$, with tile $t=4$ and tensor rank $r$ from 16 to 48 (for 49, the line is known to cross 0 by Strassen). The red line is our estimate for the corresponding error for $2:4\,$ pruning.
  • Figure 3: Mean squared error $\downarrow$ of training $\mathsf{STL}$ with Strassen based vs. random Gaussian initialization, compared for different values of $r$ before and after training. This shows "smart" initialization maintains an advantage, and is consistent with other methods we have tried.
  • Figure 4: Accuracy $\uparrow$ vs. tensor rank ($r$) for baseline T2T-Vit/7 and for the same architecture with all activation-weight MatMul s in the network trunk replaced by $\mathsf{STL}$ with tensor rank $r$. Very similar picture when including the pre-trunk (T2T) network as a followup experiment.
  • Figure 5: Comparison of singular values between the trained network's (T2T-ViT-7, see \ref{['subsec_exp_class_1']}) encoded weights and random matrices of the same sizes. The graph provides evidence that learning with $\mathsf{STL}$ occurs in a higher dimensional space.

Theorems & Definitions (8)

  • Definition 3.1: $\mathsf{STL}$
  • Remark 3.2
  • Corollary 3.3
  • Claim 3.4
  • proof
  • Example 3.5
  • Lemma A.1
  • proof