Key Query logarithmic normalization

[SergeyGubanov]

В трансформерах механизм "внимания" реализован через подсчёт суммы

$$ \sum_{n = 1}^{N} \vec{V}_{n} \exp\left( \frac{\vec{K}_{n} \vec{Q}}{\sqrt{d}} \right). $$

Как я уже писал об этом в другой теме (https://dxdy.ru/post1707690.html#p1707690) в подобной сумме для того чтобы не произошло чрезмерного усреднения по всем слагаемым нужно чтобы норма ключевых векторов была не меньше чем некий порог зависящий от логарифма от $N$:

$$ \frac{1}{\sqrt{d}} \left| \vec{K}_{n} \right|^2 > \log(N). $$

А теперь представьте что происходит в LLM трансформере при длинном контексте $N \approx 100 000$. Вот оно начало читать/писать текст, сначала $N$ маленькое, потом $N$ становится всё больше и больше. Ключи которые были вычислены в начале когда $N$ было маленьким - это одно, а ключи которые были вычислены потом когда $N$ стало большим - это совсем другое. Довольно понятно, что более поздние ключи (посчитанные когда $N$ уже стало большим) вероятнее всего будут длиннее более ранних ключей (когда $N$ было ещё маленьким). Поскольку ранние ключи короче поздних ключей, то и "внимание" к раннему тексту "уделяется" всё меньше и меньше. По идее, такой трансформер с "логарифмическим весом" забывает то, что было написано в начале. Ну, мне так кажется. Экспериментально проверить не могу. Не обладаю технической возможностью. Но если это действительно так, то довольно понятно как можно было бы спасти трансформер от логарифмического забывания контекста. Похоже, что вот такая модификация "внимания" должна была бы помочь:

$$ \sum_{n = 1}^{N} \vec{V}_{n} \exp\left( \frac{\vec{K}_{n} \vec{Q}}{\sqrt{d}} \log(N) \right). $$

Нужно просто добавить $\log(N)$ в экспоненту, тогда норма ключевых векторов не будет зависеть от длины контекста.