Геометрия Attention: почему QK Norm это не просто костыль для стабильности, а способ заставить сеть понимать смысл

Если вы следите за архитектурами современных LLM (например, LLaMA или Gemma), вы могли заметить одну маленькую, но важную деталь, которая стала стандартом де-факто – QK Norm (Query-Key Normalization).

В официальных пейперах её использование объясняют сухим математическим языком: при масштабировании моделей скалярные произведения Q * Kt  начинают неконтролируемо расти, Softmax превращается в единичный вектор (one-hot), градиенты затухают, и обучение ^[1] разваливается. Нормализация решает эту проблему численной стабильности.

Это правда. Но мне кажется, что за этим скрывается нечто гораздо более интересное. В этой статье я хочу предложить теоретический взгляд на то, почему QK Norm дает буст не только к стабильности лосса, но и к фундаментальному качеству векторных представлений (эмбеддингов) слов.

Анатомия ленивой нейросети

Давайте вспомним, что такое скалярное произведение двух векторов (Query и Key) с точки зрения ^[2] геометрии.
Формула известна со школы:

A⋅B=∣A∣×∣B∣×cos(θ)

В механизме Attention скалярное произведение определяет, насколько сильно токен A хочет смотреть на токен B. Если нейросети нужно увеличить вес внимания ^[3] между двумя словами (сделать их семантически близкими), у неё есть два пути:

1. Изменить угол: Повернуть векторы так, чтобы они указывали в одном направлении в N-мерном пространстве

2. Изменить длину:  Оставить векторы на месте, но просто удлинить их.

Нейронные сети, как известно, фантастически ленивы. Оптимизатору (тому же AdamW) гораздо проще градиентно “накачать” длину вектора (просто увеличивая веса линейного слоя), чем аккуратно вращать его в высокоразмерном пространстве, пытаясь не сломать угловые расстояния до десятков тысяч других токенов.

В итоге, без нормализации, модель учится выделять важные токены не за счет их точного позиционирования в пространстве смыслов, а за счет их «громкости». Векторы частых или грамматически доминирующих токенов раздуваются в размерах, подавляя более тонкие семантические связи.

Отрезаем путь к отступлению: Геометрия QK Norm

Что мы делаем, когда применяем RMSNorm к матрицам Query и Key перед их перемножением?

Мы принудительно делаем их длину константой

Теперь наша формула скалярного произведения схлопывается до:

Q⋅K≈cos(θ)

Скалярное произведение превращается в чистое косинусное сходство (Cosine Similarity). Все векторы токенов оказываются заперты на поверхности N-мерной гиперсферы.

И вот тут начинается магия. Мы физически отрезали нейросети возможность изменять длину вектора. У неё остался только один вариант взаимодействия с данными, изменять положение вектора (угол).

Почему это дает буст к “пониманию”

Представьте себе комнату, в которой люди пытаются договориться.
Без QK Norm (когда длина вектора свободна), если кто-то хочет, чтобы его услышали, он может просто взять рупор и начать орать. Громкость компенсирует тот факт, что он стоит в другом конце комнаты.

С применением QK Norm мы отобрали у всех рупоры. Все говорят одинаково тихо. Теперь, чтобы слова одного человека повлияли на другого, им физически необходимо подойти друг к другу.

С точки зрения машинного обучения это приводит к потрясающим структурным изменениям в латентном пространстве:

1. Смысловая плотность: Модель вынуждена филигранно кластеризовать векторы на поверхности гиперсферы. Токены собираются в строгие семантические созвездия исключительно по смыслу (по углу), а не по частоте появления в тексте.

2. Защита от доминирующих токенов: В классическом трансформере “пустые” токены (вроде знаков препинания или предлогов) часто собирают на себя огромный вес внимания просто за счет раздутой длины весов (так называемые attention sinks). На гиперсфере у точек нет массы, есть только координаты.

3. Лучшая генерализация: Когда пространство смыслов построено исключительно на относительных углах, модели проще интерполировать значения для новых, редко встречающихся слов.

Реализация на Pytorch

Код изменения классического Attention до безобразия прост. Нам нужно добавить всего пару строк:

import torch.nn as nn
import torch.nn.functional as F

class QKNormAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        
        # Добавляем нормализацию для Q и K
        self.q_norm = nn.RMSNorm(d_model // n_heads)
        self.k_norm = nn.RMSNorm(d_model // n_heads)
        
    def forward(self, x):
        # ... (стандартное разделение на головы) ...
        
        # Применяем нормализацию ДО скалярного произведения
        q = self.q_norm(q)
        k = self.k_norm(k)
        
        # Теперь это по сути косинусное расстояние
        scores = torch.matmul(q, k.transpose(-2, -1))
        # Масштабирование на sqrt(d_k) часто можно убрать или заменить на обучаемый скаляр
        # (так называемый temperature parameter)
        
        attn = F.softmax(scores, dim=-1)
        # ... (дальше как обычно) ...

Заключение

Я не проводил масштабных A/B тестов (обучение LLM с нуля дорогое удовольствие), но этот мысленный эксперимент кажется мне предельно логичным.

Индустрия пришла к QK Norm через боль ^[4]: модели масштаба 100B+ параметров просто отказывались стабильно обучаться без этого костыля. Но, возможно, решая инженерную проблему “взрыва логитов”, инженеры случайно наткнулись на мощнейший регуляризатор репрезентаций.

Лишая сеть одной степени свободы (масштаба), мы заставляем её максимально эффективно использовать оставшуюся (позиционирование), что ведет к более глубокому и геометрически выверенному пониманию языка.

А что вы думаете об этом? Замечали ли вы улучшение качества эмбеддингов при переходе на косинусное внимание в своих проектах?