paper
Transformer 和 BERT 里, 多头注意力并没有表面上看起来那么”每个头都必不可少”.
很多 head 可以删掉而几乎不损失性能, 但不同层和不同注意力类型的重要性差别很大, 尤其是编码器-解码器注意力更依赖多头结构.
对于decoder only架构, 使用128 head dim为最优解. 64dim会导致冗余的action head太多, 256 head dim会导致GPU计算的时候显存加载速度变慢. 同时, 使用128 head dim 对于 Flash Attention 而言性能最佳