HEAT：用于Transformer压缩的硬件高效自动张量分解
机器学习
2022-12-01 v1 人工智能
硬件体系结构
摘要
Transformer在自然语言处理和计算机视觉中取得了优越性能。其自注意力和前馈层存在过度参数化，限制了推理速度和能效。张量分解是一种有前景的技术，通过利用张量代数性质以因子化形式表达参数，从而减少参数冗余。先前的工作使用手动或启发式分解设置，缺乏硬件感知定制，导致硬件效率低下和较大的性能下降。本工作中，我们提出一种硬件感知的张量分解框架，称为HEAT，它能够高效探索可能分解的指数级空间，并通过硬件感知协同优化自动选择张量化形状和分解秩。我们联合研究张量收缩路径优化和融合Einsum映射策略，以弥合理论收益与实际硬件效率提升之间的差距。我们的两阶段知识蒸馏流程解决了可训练性瓶颈，从而显著提升了因子化Transformer的最终精度。总体而言，我们通过实验表明，我们的硬件感知因子化BERT变体将能量延迟积降低了5.7倍，精度损失小于1.1%，并比手工调优和启发式基线实现了更优的效率-精度帕累托边界。
引用
@article{arxiv.2211.16749,
title = {HEAT: Hardware-Efficient Automatic Tensor Decomposition for Transformer Compression},
author = {Jiaqi Gu and Ben Keller and Jean Kossaifi and Anima Anandkumar and Brucek Khailany and David Z. Pan},
journal= {arXiv preprint arXiv:2211.16749},
year = {2022}
}
备注
9 pages. Accepted to NeurIPS ML for System Workshop 2022 (Spotlight)