Efficient attention layers that use a forget gate to control what to remember are optimized for NVIDIA GPUs. On Google TPUs, which use wide matrix arrays rather than fast individual calculations, these layers run slowly because half of the computational units sit idle and vector operations become bottlenecks.
The layer is redesigned for TPUs by widening the key vectors to 256 elements, sharing each key across multiple value heads, and grouping the forget gates into four blocks. This ensures all operations use the fast matrix array.
On TPU v7x, the new layer runs 1.6 to 2.5 times faster than the best existing version, with similar accuracy. This speeds up training of large language models by up to 24 percent, making TPUs a more attractive option for AI development.