MLKV: 用于内存高效Transformer解码的多层键值头
机器学习
2024-10-16 v3
摘要
Transformer的自回归推理极大地受益于键值(KV)缓存,但随着模型规模、批大小和序列长度的大幅增长,这可能导致严重的内存瓶颈。我们引入了多层键值(MLKV)共享,这是一种新颖的方法,将KV共享扩展到Transformer的各层之间,以比多头注意力(MQA)和分组查询注意力(GQA)更大幅度地减少内存使用。在使用升级训练的Pythia-160M变体进行的各种NLP基准和推理指标评估中,MLKV显著减少了内存使用,性能损失极小,与MQA相比,KV缓存大小最多可减少6倍。这些结果凸显了MLKV在规模化高效部署Transformer模型方面的潜力。我们在https://github.com/zaydzuhri/pythia-mlkv提供代码。
引用
@article{arxiv.2406.09297,
title = {MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding},
author = {Zayd Muhammad Kawakibi Zuhri and Muhammad Farid Adilazuarda and Ayu Purwarianti and Alham Fikri Aji},
journal= {arXiv preprint arXiv:2406.09297},
year = {2024}
}