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

В классическом 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 это шаг в сторону более живых архитектур.
Мы не просто учим модель смотреть на токены,
мы учим её чувствовать движение своего внимания как будто у неё появилась инерция восприятия.