arXiv:2503.12211cs.LGcs.AI2025-03被引 7

用新算子替代矩阵乘法,提速降耗不丢精度

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

  • 用可学习的局部基变换+元素相乘替代传统矩阵乘法
  • 4×4块近似矩阵乘法,计算量减少2.66倍,参数略增
  • 适合追求算力效率的模型部署与训练场景

现代AI高度依赖大规模矩阵乘法(MatMul),其计算带来推理与训练的可扩展性瓶颈。本文提出一种面向GPU原生的双线性算子——斯特拉森-分块(Strassen-Tile, STL),在速度、精度和参数量之间实现三向权衡。该算子通过在权重与激活矩阵分块上应用可学习的局部基变换,并以元素乘积形式并行实现,评估所需浮点运算量远低于传统方法(≪ n³),但参数量高于标准MatMul(≫ n²)。关键问题是优化每层的基变换,这是一个高度非凸问题。我们发现基于理论启发的初始化(源自快速矩阵与多项式乘法)显著优于随机SGD初始化。实验表明,STL可实现4×4分块的矩阵乘法近似,使计算量降低2.66倍,同时提升SoTA T2T-ViT-7(430万参数)在ImageNet-1K上的准确率。即使使用未针对CUDA优化的PyTorch代码,仍能在计算密集型场景中获得实际运行时加速。该结果结合其理论基础,表明STL是构建可扩展、低成本AI系统的重要候选组件。

原文摘要 · Abstract (English)

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.

矩阵乘法GPU优化高效计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。