Token-Picker: Accelerating Attention in Text Generation with Minimized Memory Transfer via Probability Estimation
Abstract
The attention mechanism in text generation is memory-bounded due to its sequential characteristics. Therefore, off-chip memory accesses should be minimized for faster execution. Although previous methods addressed this by pruning unimportant tokens, they fall short in selectively removing tokens with near-zero attention probabilities in each instance. Our method estimates the probability before the softmax function, effectively removing low probability tokens and achieving an 12.1x pruning ratio without fine-tuning. Additionally, we present a hardware design supporting seamless on-demand off-chip access. Our approach shows 2.6x reduced memory accesses, leading to an average 2.3x speedup and a 2.4x energy efficiency.
Cite
@article{arxiv.2407.15131,
title = {Token-Picker: Accelerating Attention in Text Generation with Minimized Memory Transfer via Probability Estimation},
author = {Junyoung Park and Myeonggu Kang and Yunki Han and Yanggon Kim and Jaekang Shin and Lee-Sup Kim},
journal= {arXiv preprint arXiv:2407.15131},
year = {2024}
}
Comments
To appear in the proceedings of 61st Design Automation Conference (DAC)