В классическом self-attention каждый токен смотрит на другие токены, чтобы понять, что важно в данный момент.
Внимание распределяется мгновенно:

Именно этот механизм сделал трансформеры тем, чем они стали.

Но вот в чём проблема - внимание не имеет памяти.
На каждой итерации оно переобучается заново, не зная, куда оно смотрело в прошлый раз.
Из за этого внимание может скакать, шуметь и терять контекст, особенно в длинных последовательностях.

Проблема: внимание без инерции

Представьте, что вы идёте по неровной дороге.
Если вы будете менять направление мгновенно, без инерции, вас просто будет бросать из стороны в сторону.
Точно так же и внимание в трансформере:
оно то цепляется за один токен, то внезапно переключается на другой,
порождая хаотичные изменения в градиентах и мешая стабильному обучению.

А что, если добавить вниманию немного физики?

Momentum это понятие из механики.
Если у тела есть скорость, оно не останавливается мгновенно, а плавно замедляется.
Почему бы не применить тот же принцип к вниманию?

Идея:

Пусть текущее внимание немного зависит от того, каким оно было раньше.
Не только “куда я смотрю сейчас?”,
но и “куда я смотрел мгновение назад?”.

От классического внимания к Momentum Attention

В классике:

Теперь добавим инерцию к Value-векторам:

_{Пояснение: Если бы я добавил инерцию к attn_scores, модель была бы вынуждена смотреть на те же самые токены, что и на прошлом шаге. Это очень жесткое ограничение. Добавляя инерцию к V, я позволяю вниманию свободно выбирать, куда смотреть на каждом шаге (Q и K новые), но информация, которую оно извлекает (V), будет смесью новой и старой.}

Тогда:

То есть текущее внимание теперь частично помнит, какие значения были важны на предыдущем шаге. α (например, 0.9) задаёт вес настоящего по сравнению с прошлым

Простой пример на pytorch

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

class MomentumAttention(nn.Module):
    def __init__(self, d_model, n_heads=8, alpha=0.9):
        super().__init__()
        if d_model % n_heads != 0:
            raise ValueError("d_model должен делиться на n_heads без остатка")
            
        self.alpha = alpha
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, Q, K, V, prev_V=None):

        B, T_q, D = Q.shape
        _, T_k, _ = K.shape

        # Линейные проекции и разделение на головы
        q = self.W_q(Q).view(B, T_q, self.n_heads, self.d_k).transpose(1, 2) # [B, n_heads, T_q, d_k]
        k = self.W_k(K).view(B, T_k, self.n_heads, self.d_k).transpose(1, 2) # [B, n_heads, T_k, d_k]
        v = self.W_v(V).view(B, T_k, self.n_heads, self.d_k).transpose(1, 2) # [B, n_heads, T_k, d_k]

        # Применение Momentum к векторам Value
        if prev_V is None:
            # На самом первом шаге инерции нет, используем текущее значение
            v_momentum = v
        else:
            # Совмещаем текущее значение с прошлым
            v_momentum = self.alpha * v + (1 - self.alpha) * prev_V

        # 3. Сохраняем новое состояние для следующего шага.
        # .detach() используется, чтобы градиенты не текли через всю историю состояний,
        # что превратило бы механизм в полноценный RNN и сильно усложнило бы обучение.
        new_prev_V = v_momentum.detach()

        # 4. Стандартный механизм self-attention
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5)
        attn_weights = F.softmax(attn_scores, dim=-1)
        
        # Внимание применяется к инерционным значениям v_momentum
        out = torch.matmul(attn_weights, v_momentum)

        # 5. Собираем головы вместе и пропускаем через финальный линейный слой
        out = out.transpose(1, 2).contiguous().view(B, T_q, D)
        
        return self.W_o(out), new_prev_V


# Пример модели, которая использует MomentumAttention

class AutoregressiveModel(nn.Module):
    def __init__(self, vocab_size, d_model, n_heads, alpha):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.momentum_attn = MomentumAttention(d_model, n_heads, alpha)
        self.layernorm = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.ReLU(),
            nn.Linear(4 * d_model, d_model)
        )
        self.out_proj = nn.Linear(d_model, vocab_size)

    def forward(self, input_ids):

        B, T = input_ids.shape
        x = self.embedding(input_ids)

        # Инициализируем состояние для всей последовательности
        prev_V_state = None
        all_step_outputs = []

        # Цикл по каждому шагу (токену) в последовательности
        for t in range(T):
            # Берем срез данных для текущего шага
            # В реальном декодере Q - это текущий токен, K и V - все предыдущие.
            # Для простоты демонстрации механизма инерции, мы используем только текущий токен
            # как Q, K, и V. Это показывает, как состояние `prev_V_state` передается.
            current_x_step = x[:, t:t+1, :] # Shape: [B, 1, D]
            
            # Вызываем слой внимания, передавая ему состояние с прошлого шага
            attn_output, prev_V_state = self.momentum_attn(
                Q=current_x_step, 
                K=current_x_step, 
                V=current_x_step, 
                prev_V=prev_V_state
            )
            
            # Стандартные блоки трансформера (residual connection, layernorm, FFN)
            h = self.layernorm(current_x_step + attn_output)
            step_output = self.ffn(h)
            all_step_outputs.append(step_output)

        # Собираем выходы со всех шагов в один тензор
        full_output = torch.cat(all_step_outputs, dim=1) # Shape: [B, T, D]
        
        # Финальная проекция в размер словаря
        logits = self.out_proj(full_output)
        return logits


# Параметры
batch_size = 4
seq_len = 10
vocab_size = 100
d_model = 64
n_heads = 8
alpha = 0.9

# Создаем модель
model = AutoregressiveModel(vocab_size, d_model, n_heads, alpha)

# Создаем случайные входные данные
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))

print(f"Входные данные (shape): {input_ids.shape}")

# Получаем выход модели
output_logits = model(input_ids)

print(f"Выходные логиты (shape): {output_logits.shape}")

# Проверка корректности размеров
assert output_logits.shape == (batch_size, seq_len, vocab_size)

print("\nМодель успешно отработала")

Что это даёт

Сглаживание представлений.
Вектора V не перескакивают резко между шагами прошлое состояние частично сохраняется, что снижает турбулентность активаций.
Более стабильное распределение внимания.
Модель получает эффект инерции в значениях, и внимание не скачет при малых изменениях входа. Это особенно полезно в авторегрессионных моделях, где выходы сильно зависят от предыдущего шага.
Облегчённое обучение.
Так как prev_V передаётся через detach(), градиенты не текут сквозь всю историю, что предотвращает взрыв или затухание градиентов в отличие от полного RNN-подхода.
Простая интеграция.
Механизм не требует изменения архитектуры он полностью совместим с обычным MultiHeadAttention и может быть вставлен в любой трансформерный блок.

Возможные минусы

Накопление смещения (drift).
Если alpha слишком велико, старые состояния начинают тянуть новые векторные представления, и внимание может начать запоминать шум.
Сложность выбора alpha.
Значение 0.9 подходит не всегда при быстрых изменениях контекста модель может терять реактивность (поздно реагировать на новые токены).
Невозможность параллелизации по времени.
Так как состояние prev_V передаётся последовательно, обучение по всей последовательности становится менее параллельным (особенно при autoregressive setup).
Потенциальная инерция ошибок.
Если модель делает ошибку на шаге t, она может частично переноситься дальше через prev_V, особенно при большом alpha.

Заключение

Momentum Attention это шаг в сторону более живых архитектур.
Мы не просто учим модель смотреть на токены,
мы учим её чувствовать движение своего внимания как будто у неё появилась инерция восприятия.

Momentum Attention: когда внимание получает инерцию