arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2406.06248cs.LG
收藏

更优的计算投入:用结构化矩阵替代稠密层

Compute Better Spent: Replacing Dense Layers with Structured Matrices

New York University(纽约大学);Carnegie Mellon University(卡内基梅隆大学)

Shikai Qiu, Andres Potapczynski, Marc Finzi, Micah Goldblum, Andrew Gordon Wilson

AI总结:

本文系统研究用结构化矩阵替代稠密层,提出包含 Monarch 矩阵的块张量训练(BTT)矩阵族,通过最大更新参数化确定最优初始化与学习率缩放,在相同计算量下于图像分类和语言建模任务上超越稠密矩阵。

摘要:

稠密线性层是基础模型中的主要计算瓶颈。寻找比稠密矩阵更高效的替代方案,对于构建计算效率更高的模型具有巨大潜力,卷积网络在图像领域的成功即为明证。在本工作中,我们系统地探索将结构化矩阵作为稠密矩阵的替代品。我们表明,不同的结构通常需要截然不同的初始化尺度和学习率,这对性能至关重要,尤其是在模型规模扩大时。利用最大更新参数化(Maximal Update Parameterization)的见解,我们确定了这些非常规层的最优初始化与学习率缩放。最后,我们测量了不同结构的缩放规律,以比较其性能随计算量提升的速度。我们提出了一个包含 Monarch 矩阵的新型矩阵族——块张量训练(Block Tensor-Train,BTT),并证明在多个任务上,在相同计算量下 BTT 的性能优于稠密矩阵。在带数据增强的 CIFAR-10/100 上,训练 MLP 和 ViT 时,BTT 的训练损失以指数级低于稠密矩阵。BTT 在 ImageNet-1k 上以少 3.8 倍的计算量达到与稠密 ViT-S/32 相当的性能,并且在训练小型 GPT-2 语言模型时比稠密矩阵更高效。

英文摘要:

Dense linear layers are the dominant computational bottleneck in foundation models. Identifying more efficient alternatives to dense matrices has enormous potential for building more compute-efficient models, as exemplified by the success of convolutional networks in the image domain. In this work, we systematically explore structured matrices as replacements for dense matrices. We show that different structures often require drastically different initialization scales and learning rates, which are crucial to performance, especially as models scale. Using insights from the Maximal Update Parameterization, we determine the optimal scaling for initialization and learning rates of these unconventional layers. Finally, we measure the scaling laws of different structures to compare how quickly their performance improves with compute. We propose a novel matrix family containing Monarch matrices, the Block Tensor-Train (BTT), which we show performs better than dense matrices for the same compute on multiple tasks. On CIFAR-10/100 with augmentation, BTT achieves exponentially lower training loss than dense when training MLPs and ViTs. BTT matches dense ViT-S/32 performance on ImageNet-1k with 3.8 times less compute and is more efficient than dense for training small GPT-2 language models.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑