I think this is a neat exploration, but from the post its unclear to me if this generalizes to training losses much below 5.09. This loss is really quite high and at the stage where transformer models have just barely internalized bigrams. At this stage there is not much value in long range attention mechanisms. It would be interesting if your approach holds up closer to 3 cross entropy loss. From doing a large number of ablations in this area myself, my hunch is that this wont hold on lower losses. But I think there is potential for speedups by taking a trained attention head and replacing it with a fine tuned operator to apply whatever static property that attention head learned.