[论文解读] Kaleidoscope: An Efficient, Learnable Representation For All Structured Linear Maps

Tri Dao, Nimit S. Sohoni+6 more|arXiv (Cornell University)|Dec 29, 2020

Neural Networks and Applications参考文献 79被引用 9

一句话总结

K-matrices（K-矩阵）是一种可学习的、参数高效的结构化线性映射表示方法，通过组合蝴蝶矩阵来捕捉低秩、稀疏和傅里叶变换等多种结构，实现接近最优的空间和时间复杂度。该方法支持端到端可微学习，在ImageNet上将模型准确率提升最高达5%，在Transformer中将推理速度提升36%。

ABSTRACT

Modern neural network architectures use structured linear transformations, such as low-rank matrices, sparse matrices, permutations, and the Fourier transform, to improve inference speed and reduce memory usage compared to general linear maps. However, choosing which of the myriad structured transformations to use (and its associated parameterization) is a laborious task that requires trading off speed, space, and accuracy. We consider a different approach: we introduce a family of matrices called kaleidoscope matrices (K-matrices) that provably capture any structured matrix with near-optimal space (parameter) and time (arithmetic operation) complexity. We empirically validate that K-matrices can be automatically learned within end-to-end pipelines to replace hand-crafted procedures, in order to improve model quality. For example, replacing channel shuffles in ShuffleNet improves classification accuracy on ImageNet by up to 5%. K-matrices can also simplify hand-engineered pipelines -- we replace filter bank feature computation in speech data preprocessing with a learnable kaleidoscope layer, resulting in only 0.4% loss in accuracy on the TIMIT speech recognition task. In addition, K-matrices can capture latent structure in models: for a challenging permuted image classification task, a K-matrix based representation of permutations is able to learn the right latent structure and improves accuracy of a downstream convolutional model by over 9%. We provide a practically efficient implementation of our approach, and use K-matrices in a Transformer network to attain 36% faster end-to-end inference speed on a language translation task.

研究动机与目标

• 解决在深度学习中为效率和准确率而手动选择和设计结构化线性映射（如低秩、稀疏、傅里叶变换）的挑战。
• 克服现有结构化矩阵类别的局限性，这些类别要么不可微，要么效率低下，或对多样化结构的表达能力不足。
• 开发一种统一、可微且高效的表示方法，以近乎最优的参数量和算术复杂度捕捉所有结构化线性映射。
• 实现神经网络中结构化组件的端到端学习，替代手工设计的流程（如通道混洗或滤波器组）。
• 在实际任务中展示K-矩阵的实用性，包括图像分类、语音识别和Transformer推理。

提出的方法

• 提出K-矩阵（K-matrices）作为结构化蝴蝶矩阵的乘积，后者以高效的矩阵-向量乘法著称。
• 利用理论基础：任何具有 $ s \ll n^2 $ 次操作的结构化线性映射，均可表示为总参数量匹配其算术电路复杂度的稀疏矩阵乘积。
• 在蝴蝶矩阵中使用预定义的、固定的稀疏模式，避免在稀疏分解中对稀疏模式进行不可微的搜索。
• 通过使整个K-矩阵参数化过程可微，实现端到端反向传播，支持基于梯度的优化。
• 设计一种利用蝴蝶结构的高效矩阵-向量乘法算法，实现 $ O(n \log n) $ 的时间复杂度。
• 将K-矩阵作为可学习层集成到神经网络中，包括Transformer，以替代固定的结构化组件（如通道混洗或FFT层）。

实验结果

研究问题

• RQ1 单一、统一的可微参数化方法是否能够以近乎最优的效率表示所有主要类别的结构化线性映射？
• RQ2 K-矩阵是否能在端到端深度学习流程中被有效学习，从而相比手工设计的结构化组件提升模型性能？
• RQ3 在语音识别中，K-矩阵在多大程度上可以替代传统预处理步骤（如滤波器组计算），同时保持极低的准确率损失？
• RQ4 K-矩阵是否能够发现复杂任务（如置换图像分类）中的潜在结构，并提升下游模型的准确率？
• RQ5 在真实模型（如Transformer）中，用K-矩阵替换标准线性层后，实际推理速度提升和内存效率增益如何？

主要发现

• K-矩阵能够以接近最优的参数量和算术复杂度表示任意结构化线性映射，其复杂度与矩阵算术电路复杂度的理论下界一致。
• 在ShuffleNet中用可学习的K-矩阵层替代通道混洗，可在不增加推理成本的前提下，将ImageNet的top-1准确率提升最高达5%。
• 在TIMIT语音识别任务中，用可学习的K-矩阵层替代手工设计的滤波器组，仅导致0.4%的准确率下降，同时显著简化了模型流程。
• 在一项具有挑战性的置换图像分类任务中，基于K-矩阵的置换表示使下游卷积模型的准确率提升超过9%。
• 将K-矩阵集成到Transformer模型中，在语言翻译任务上实现了36%的端到端推理速度提升，且准确率无显著损失。
• K-矩阵表示支持结构化组件的高效、可微学习，消除了训练过程中对启发式或高成本稀疏模式搜索的需求。