Механизм внимания в PyTorch

Механизм внимания — это одна из самых важных концепций в глубоком обучении.

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

В этом разделе подробно описаны основные принципы механизма внимания, реализации на PyTorch и различные варианты внимания.

Применимая версия:Код в этой статье написан на основе PyTorch 2.0+.nn.MultiheadAttentionизbatch_firstПараметр был введён в PyTorch 1.9; в более ранних версиях требуется ручная обработка размерностей.

1. Основы механизма внимания

1.1 Зачем нужно внимание

Традиционные модели последовательность-в-последовательность (Seq2Seq) имеют фундаментальную проблему:Энкодеру необходимо сжать всю информацию в вектор фиксированной длины.Для длинных последовательностей этот вектор становится информационным узким местом — чем длиннее предложение, тем серьёзнее потеря информации.

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

1.2 Суть механизма внимания

Механизм внимания можно рассматривать каквзвешенную сумму(операцию). При наличии запроса (Query), ключа (Key) и значения (Value) веса назначаются путём вычисления сходства между Query и каждым Key, после чего вычисляется взвешенная сумма Values:

\[\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\]$$

где:

Q（Query）: вектор запроса, представляющий «какую информацию я ищу»
K（Key）: вектор ключа, указывающий, «какая информация у меня здесь есть» (используется для сопоставления)
V（Value）: вектор значения, представляющий «что я могу фактически предоставить»
$$d_k$$: размерность Key, $$\sqrt{d_k}$$ используется для масштабирования, чтобы значения скалярного произведения не были слишком большими, что привело бы к затуханию градиентов softmax

Зачем делить на $$\sqrt{d_k}$$? Когда $$d_k$$ велико, дисперсия скалярного произведения Q и K растёт линейно с размерностью, из-за чего входные значения softmax становятся слишком большими, а градиенты приближаются к нулю. После масштабирования дисперсия восстанавливается до 1, обеспечивая нормальный поток градиентов.

Пример

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

def scaled_dot_product_attention(Q, K, V, mask=None):
"""
Масштабированное внимание со скалярным произведением

Параметры:
Q: тензор запроса [batch, n_heads, seq_len_q, d_k]
K: тензор ключа [batch, n_heads, seq_len_k, d_k]
V: Тензор значений [batch, n_heads, seq_len_v, d_v]
mask: Тензор маски; маскированные позиции устанавливаются в False/0

Возвращает:
output: Выход внимания [batch, n_heads, seq_len_q, d_v]
attention_weights: Веса внимания [batch, n_heads, seq_len_q, seq_len_k]
"""
d_k = Q.size(-1)

# 1. Вычислить QK^T / sqrt(d_k)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

# 2. Применить маску (необязательно)
# Установить маскированные позиции в очень маленькое значение; после softmax они приближаются к 0
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)

# 3. Нормализация softmax → веса внимания
attention_weights = F.softmax(scores, dim=-1)

# 4. Взвешенная сумма
output = torch.matmul(attention_weights, V)

return output, attention_weights

#
batch_size, n_heads = 2, 4
seq_len_q, seq_len_k = 5, 6
d_k, d_v = 8, 16

Q = torch.randn(batch_size, n_heads, seq_len_q, d_k)
K = torch.randn(batch_size, n_heads, seq_len_k, d_k)
V = torch.randn(batch_size, n_heads, seq_len_k, d_v)

output, attn_weights = scaled_dot_product_attention(Q, K, V)

print(f"Форма вывода: {output.shape}") # [2, 4, 5, 16]
print(f"Форма весов внимания: {attn_weights.shape}") # [2, 4, 5, 6]
print(f"Сумма по строке весов (должна быть 1.0): {attn_weights[0, 0, 0].sum().item():.4f}")

Суть весов внимания — это распределение вероятностей: после softmax сумма каждой строки равна 1. Чем больше вес, тем сильнее модель "фокусируется" на этой позиции.

2. Модуль внимания PyTorch

2.1 Многоголовое внимание

Многоголовое вниманиеОно позволяет модели одновременно обращать внимание на информацию из разных подпространств представлений в разных позициях. Оно проецирует Q, K и V с помощью разных линейных проекций в несколько подпространств, независимо вычисляет внимание в каждом подпространстве и в конце конкатенирует результаты и применяет ещё одно линейное преобразование.

Это как если бы несколько человек одновременно рассматривали один и тот же текст с разных точек зрения — одни сосредоточены на грамматической структуре, другие на семантических связях, а третьи на дальних зависимостях. Механизм многоголового внимания позволяет модели улавливать более богатые закономерности.

$$\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O$$

$$\text{где head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

Пример

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

class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0, "d_model должно делиться на n_heads"

self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads # Размерность каждой головы

# Линейные проекции для Q, K, V (вычислить все головы сразу, более эффективно)
self.w_q = nn.Linear(d_model, d_model)
self.w_k = nn.Linear(d_model, d_model)
self.w_v = nn.Linear(d_model, d_model)

# Выходная проекция
self.w_o = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)

def forward(self, query, key, value, mask=None):
"""
Параметры:
query: [batch, seq_len_q, d_model]
key: [batch, seq_len_k, d_model]
value: [batch, seq_len_k, d_model]
mask: [batch, 1, 1, seq_len_k] или [batch, 1, seq_len_q, seq_len_k]
"""
batch_size = query.size(0)

# 1. Линейная проекция, затем разделение на головы
# [batch, seq_len, d_model] → [batch, seq_len, n_heads, d_k]
# → [batch, n_heads, seq_len, d_k]
Q = self.w_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

# 2. Масштабированное внимание на основе скалярного произведения
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)

attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)

# 3. Взвешенная сумма
context = torch.matmul(attn_weights, V)

# 4. Объединение голов: [batch, n_heads, seq_len, d_k] → [batch, seq_len, d_model]
context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

# 5. Выходная проекция
output = self.w_o(context)

return output, attn_weights

#
d_model, n_heads = 128, 8
seq_len, batch = 10, 4

layer = MultiHeadAttention(d_model, n_heads)

query = torch.randn(batch, seq_len, d_model)
key = torch.randn(batch, seq_len, d_model)
value = torch.randn(batch, seq_len, d_model)

output, attn_weights = layer(query, key, value)

print(f"Форма вывода: {output.shape}") # [4, 10, 128]
print(f"Форма весов внимания: {attn_weights.shape}") # [4, 8, 10, 10]
print(f"Размерность каждой головы d_k = {d_model // n_heads}")

2.2 Встроенный MultiheadAttention в PyTorch

PyTorch предоставляет высокооптимизированныйnn.MultiheadAttention, который использует объединённые ядра на внутреннем уровне и в большинстве сценариев быстрее ручных реализаций. Рекомендуется использовать в производственных средах.

Пример

import torch
import torch.nn as nn

class TransformerAttention(nn.Module):
"""Слой самовнимания, использующий встроенный MultiheadAttention из PyTorch"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
self.attention = nn.MultiheadAttention(
embed_dim=d_model,
num_heads=n_heads,
dropout=dropout,
batch_first=True, # Формат ввода: [batch, seq, features]
# В PyTorch 2.0+ можно включить бэкенд Flash Attention:
# attn_implementation="flash_attention_2" # требуется установить flash-attn
)
self.layernorm = nn.LayerNorm(d_model)

def forward(self, x, key_padding_mask=None):
"""
Самовнимание: Q, K, V — все равны x
Параметры:
x: [batch, seq_len, d_model]
key_padding_mask: [batch, seq_len], True указывает, что эта позиция — padding
"""
attn_output, attn_weights = self.attention(
x, x, x, # самовнимание
key_padding_mask=key_padding_mask
)

# Остаточное соединение с Pre-Norm (обучение стабильнее, чем с Post-Norm)
output = self.layernorm(x + attn_output)

return output, attn_weights

#
d_model, n_heads = 128, 8
seq_len, batch = 10, 4

model = TransformerAttention(d_model, n_heads)
x = torch.randn(batch, seq_len, d_model)

# Создать маску паддинга: True указывает, что позицию следует игнорировать
key_padding_mask = torch.zeros(batch, seq_len, dtype=torch.bool)
key_padding_mask[0, 7:] = True # Последние 3 позиции первого образца являются паддингом
key_padding_mask[1, 5:] = True # Последние 5 позиций второго образца являются паддингом

output, attn_weights = model(x, key_padding_mask=key_padding_mask)

print(f"Форма вывода: {output.shape}") # [4, 10, 128]
print(f"Форма весов внимания: {attn_weights.shape}") # [4, 10, 10]

Обратите внимание на типы масок:nn.MultiheadAttentionизkey_padding_maskСоглашениеTrue = ignoreиспользуется, противоположно соглашению в некоторых пользовательских реализациях, где 0 = ignore. Обязательно проверяйте при использовании.

3. Варианты механизма внимания

3.1 Самовнимание

Самовнимание— это особая форма механизма внимания: Q, K, V все поступают из одной и той же входной последовательности. Оно позволяет каждой позиции в последовательности напрямую взаимодействовать со всеми другими позициями, тем самым захватывая глобальные зависимости. Это основной компонент Transformer и ключ к его преимуществу над RNN: RNN нужно передавать информацию шаг за шагом, а самовнимание делает это за один шаг.

Пример

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

class SelfAttention(nn.Module):
"""Слой многоголового самовнимания, поддерживает причинную маску (для авторегрессионной генерации)"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads

self.w_q = nn.Linear(d_model, d_model)
self.w_k = nn.Linear(d_model, d_model)
self.w_v = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)

def _split_heads(self, x, batch_size):
"""[batch, seq, d_model] → [batch, n_heads, seq, d_k]"""
return x.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

def forward(self, x, causal=False):
"""
Самовнимание: Q, K, V все поступают из одного и того же входа x
Параметры:
x: [batch, seq_len, d_model]
causal: использовать ли причинную маску (чтобы не видеть будущие позиции)
"""
batch_size, seq_len, _ = x.size()

# Проекция + разделение на головы
Q = self._split_heads(self.w_q(x), batch_size)
K = self._split_heads(self.w_k(x), batch_size)
V = self._split_heads(self.w_v(x), batch_size)

# Вычислить оценки внимания
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

# Каузальная маска: установить верхнюю треугольную матрицу в -inf, чтобы предотвратить просмотр будущих токенов
if causal:
causal_mask = torch.triu(
torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool),
diagonal=1
)
scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))

attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)

# Взвешенная сумма + объединение голов
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)

return self.out_proj(context), attn_weights

#
d_model, n_heads = 256, 8
seq_len, batch = 20, 2

attn = SelfAttention(d_model, n_heads)
x = torch.randn(batch, seq_len, d_model)

# Без каузальной маски (сценарий энкодера)
output_enc, weights_enc = attn(x, causal=False)
print(f"[Энкодер] Выход: {output_enc.shape}, веса: {weights_enc.shape}")

# С каузальной маской (сценарий декодера / авторегрессивной генерации)
output_dec, weights_dec = attn(x, causal=True)
print(f"[Декодер] Выход: {output_dec.shape}, веса: {weights_dec.shape}")

# Проверка каузальности: первый токен не должен обращать внимание на последующие токены
print(f"causal=True, weights[0,0,0,5:] (должны быть все 0): {weights_dec[0, 0, 0, 5:].tolist()[:5]}")

3.2 Перекрёстное внимание

В кросс-внимании, Q приходит из одной последовательности (например, Decoder), а K и V — из другой последовательности (например, Encoder). Это мост архитектуры Encoder-Decoder — он позволяет Decoder динамически «искать» все выходы Encoder при генерации каждого токена.

Пример

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

class CrossAttention(nn.Module):
"""Кросс-внимание: Q приходит из целевой последовательности, K/V — из исходной последовательности"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads

self.w_q = nn.Linear(d_model, d_model)
self.w_k = nn.Linear(d_model, d_model)
self.w_v = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)

def forward(self, target, source, mask=Отсутствует):
"""
Параметры:
target: целевая последовательность [batch, target_len, d_model] → предоставляет Q
source: исходная последовательность [batch, source_len, d_model] → предоставляет K, V
mask: необязательная маска
"""
batch_size = target.size(0)
target_len = target.size(1)
source_len = source.size(1)

# Проекция + разделение на головы
Q = self.w_q(target).view(batch_size, target_len, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(source).view(batch_size, source_len, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(source).view(batch_size, source_len, self.n_heads, self.d_k).transpose(1, 2)

# Вычисление внимания
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))

attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)

# Взвешенная сумма + объединение голов
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, target_len, self.d_model)

return self.out_proj(context), attn_weights

#
d_model, n_heads = 256, 8
batch = 2
target_len, source_len = 10, 15

cross_attn = CrossAttention(d_model, n_heads)
target = torch.randn(batch, target_len, d_model) # Выход декодера
source = torch.randn(batch, source_len, d_model) # Выход энкодера

output, weights = cross_attn(target, source)

print(f"Форма целевой последовательности: {target.shape}")
print(f"Форма исходной последовательности: {source.shape}")
print(f"Форма выхода: {output.shape}") # [2, 10, 256]
print(f"Форма весов: {weights.shape}") # [2, 8, 10, 15]
# Смысл весов: насколько каждая позиция цели обращает внимание на каждую позицию источника

Самовнимание vs. перекрёстное внимание:В самовнимании Q/K/V имеют одинаковую форму ($$seq \times seq$$ матрица весов); в перекрёстном внимании длина Q может отличаться от длины K ($$target \times source$$ матрица весов).

3.3 Позиционное кодирование

Сам механизм внимания обладаетинвариантностью к перестановке—перемешивание порядка входа не меняет выход. Но язык и изображения упорядочены, поэтому позиционную информацию необходимо явно вводить через позиционное кодирование. Оригинальный Transformer использует функции синуса/косинуса для генерации фиксированных позиционных кодировок:

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right), \quad PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$

Пример

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
"""Синусоидальное позиционное кодирование (оригинальная схема Transformer)"""
def __init__(self, d_model, max_len=5000, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(p=dropout)

# Предварительно вычислить матрицу позиционного кодирования [max_len, d_model]
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

# div_term: управляет скоростью затухания разных частотных размерностей
div_term = torch.exp(
torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
)

pe[:, 0::2] = torch.sin(position * div_term) # Чётные размерности
pe[:, 1::2] = torch.cos(position * div_term) # Нечётные размерности

pe = pe.unsqueeze(0) # [1, max_len, d_model]
self.register_buffer('pe', pe) # Не обучается

def forward(self, x):
"""
x: [batch, seq_len, d_model]
Возвращает: x + positional_encoding
"""
x = x + self.pe[:, :x.size(1), :]
return self.dropout(x)

#
d_model = 128
seq_len = 50
batch = 4

pe = PositionalEncoding(d_model)
x = torch.randn(batch, seq_len, d_model)
output = pe(x)

print(f"Форма входа: {x.shape}")
print(f"Форма вывода: {output.shape}")
print(f"Форма позиционного кодирования: {pe.pe.shape}") # [1, 5000, 128]

4. Визуальные применения механизма внимания

Механизмы внимания применимы не только к последовательным данным, но также успешно используются в компьютерном зрении. Далее представлены три классических модуля визуального внимания:

4.1 Канальное внимание

Канальное внимание(например, SENet) позволяет сети учиться, "какие каналы более важны". Оно сжимает пространственную информацию с помощью глобального пулинга по каждому каналу, затем изучает межканальные зависимости через два полносвязных слоя.

Пример

import torch
import torch.nn as nn

class SEBlock(nn.Module):
"""Блок Squeeze-and-Excitation (канальное внимание)"""
def __init__(self, channels, reduction=16):
super().__init__()
self.squeeze = nn.AdaptiveAvgPool2d(1) # Глобальный средний пулинг
self.excitation = nn.Sequential(
nn.Linear(channels, channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(channels // reduction, channels, bias=False),
nn.Sigmoid() # Веса выходных каналов от 0 до 1
)

def forward(self, x):
# x: [batch, channels, H, W]
b, c, _, _ = x.size()

# Squeeze: [B, C, H, W] → [B, C, 1, 1] → [B, C]
y = self.squeeze(x).view(b, c)

# Excitation: изучение весов каналов [B, C]
y = self.excitation(y).view(b, c, 1, 1)

# Взвешивание каналов: умножить каждый канал на соответствующий ему вес
return x * y.expand_as(x)

#
se = SEBlock(channels=256, reduction=16)
x = torch.randn(4, 256, 32, 32)
output = se(x)
print(f"Вход: {x.shape} → Выход: {output.shape}")

4.2 Пространственное внимание

Пространственное вниманиепозволяет сети определять, «какие позиции более важны». Оно сжимает по размерности каналов (беря среднее и максимум), затем использует свёртку, чтобы выучить карту пространственных весов.

Пример

import torch
import torch.nn as nn

class SpatialAttention(nn.Module):
"""Модуль пространственного внимания"""
def __init__(self, kernel_size=7):
super().__init__()
padding = kernel_size // 2
# На входе 2 канала (среднее + максимум), на выходе 1 канал (карта внимания)
self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
self.sigmoid = nn.Sigmoid()

def forward(self, x):
# x: [batch, channels, H, W]

# Сжатие каналов: взять среднее и максимум → по одному [B, 1, H, W] для каждого
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)

# Конкатенация → [B, 2, H, W]
scale = torch.cat([avg_out, max_out], dim=1)

# Свёртка + Sigmoid → карта пространственных весов [B, 1, H, W]
scale = self.sigmoid(self.conv(scale))

# Пространственное взвешивание
return x * scale

#
spatial_attn = SpatialAttention(kernel_size=7)
x = torch.randn(4, 256, 32, 32)
output = spatial_attn(x)
print(f"Вход: {x.shape} → Выход: {output.shape}")

4.3 CBAM (свёрточный блок внимания)

CBAM применяет внимание к каналам и пространственное внимание впоследовательном режиме: сначала выбирает важные каналы с помощью внимания к каналам, затем выбирает важные позиции с помощью пространственного внимания. Эти два механизма дополняют друг друга, и эффект лучше, чем при использовании любого из них по отдельности.

Пример

import torch
import torch.nn as nn

class CBAM(nn.Module):
"""CBAM: внимание к каналам + пространственное внимание"""
def __init__(self, channels, reduction=16, kernel_size=7):
super().__init__()
# Внимание к каналам
self.channel_attn = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(channels, channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(channels // reduction, channels, bias=False),
nn.Sigmoid()
)
# Пространственное внимание
self.spatial_attn = nn.Sequential(
nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False),
nn.Sigmoid()
)

def forward(self, x):
# 1. Внимание к каналам
b, c, _, _ = x.size()
ca = self.channel_attn(x).view(b, c, 1, 1)
x = x * ca

# 2. Пространственное внимание
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
sa = self.spatial_attn(torch.cat([avg_out, max_out], dim=1))
x = x * sa

return x

#
cbam = CBAM(channels=256)
x = torch.randn(4, 256, 32, 32)
output = cbam(x)
print(f"Вход: {x.shape} → Выход: {output.shape}")

5. Применения механизма внимания в NLP

5.1 Кодировщик Transformer

Кодировщик Transformer состоит из N одинаковых слоёв, сложенных вместе, каждый из которых содержит два подуровня:многоголовое самовниманиеисеть прямого распространения. Каждый подуровень использует остаточные соединения и нормализацию слоя.

Пример

import torch
import torch.nn as nn
import math

class TransformerEncoderLayer(nn.Module):
"""Один слой кодировщика Transformer"""
def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
super().__init__()
# Многоголовое самовнимание
self.self_attn = nn.MultiheadAttention(
d_model, n_heads, dropout=dropout, batch_first=True
)
# Сеть прямого распространения (два линейных слоя + активация)
self.ffn = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(), # GELU используется чаще, чем ReLU
nn.Dropout(dropout),
nn.Linear(d_ff, d_model)
)
# Нормализация слоя (стиль Pre-Norm, более стабильное обучение)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)

def forward(self, x, key_padding_mask=None):
# Подслой 1: Самовнимание
residual = x
x = self.norm1(x) # Pre-Norm
attn_out, _ = self.self_attn(x, x, x, key_padding_mask=key_padding_mask)
x = residual + self.dropout(attn_out)

# Подслой 2: Сеть прямого распространения
residual = x
x = self.norm2(x) # Pre-Norm
x = residual + self.dropout(self.ffn(x))

return x

class TransformerEncoder(nn.Module):
"""Многослойный кодировщик Transformer"""
def __init__(self, d_model, n_heads, d_ff, num_layers, dropout=0.1):
super().__init__()
self.layers = nn.ModuleList([
TransformerEncoderLayer(d_model, n_heads, d_ff, dropout)
for _ in range(num_layers)
])
self.final_norm = nn.LayerNorm(d_model)

def forward(self, x, key_padding_mask=None):
for layer in self.layers:
x = layer(x, key_padding_mask)
return self.final_norm(x)

#
d_model, n_heads, d_ff = 128, 4, 512
num_layers, seq_len, batch = 3, 20, 4

encoder = TransformerEncoder(d_model, n_heads, d_ff, num_layers)
x = torch.randn(batch, seq_len, d_model)

# Маска padding: True = игнорировать эту позицию
key_padding_mask = torch.zeros(batch, seq_len, dtype=torch.bool)
key_padding_mask[0, 15:] = True

output = encoder(x, key_padding_mask=key_padding_mask)
print(f"Вход: {x.shape} → Выход: {output.shape}")

5.2 Визуализация внимания

Визуализация весов внимания — важное средство понимания и отладки моделей Transformer. С помощью тепловых карт можно интуитивно увидеть, на какую информацию модель обращает внимание в каждой позиции.

Пример

import torch
import matplotlib.pyplot as plt
import numpy as np

def visualize_attention(attention_weights, tokens=None, save_path=Отсутствует):
"""
Визуализация весов многоголового внимания

Параметры:
attention_weights: [n_heads, seq_len, seq_len]
tokens: список токенов (необязательно)
save_path: путь сохранения (необязательно)
"""
n_heads = attention_weights.shape[0]
seq_len = attention_weights.shape[1]

fig, axes = plt.subplots(1, n_heads, figsize=(n_heads * 3, 3.5))
if n_heads == 1:
axes = [axes]

for head in range(n_heads):
ax = axes[head]
attn = attention_weights[head].detach().cpu().numpy()

im = ax.imshow(attn, cmap='Blues', aspect='auto', vmin=0, vmax=attn.max())
ax.set_title(f'Голова {head + 1}', fontsize=10, fontweight='bold')
ax.set_xlabel('Ключ')
ax.set_ylabel('Запрос')

if tokens is not None:
ax.set_xticks(range(len(tokens)))
ax.set_yticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=90, fontsize=8)
ax.set_yticklabels(tokens, fontsize=8)

plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.show()

# Смоделированные веса внимания после обучения
n_heads, seq_len = 4, 8
attention_weights = torch.softmax(torch.randn(n_heads, seq_len, seq_len), dim=-1)
tokens = ['The', 'cat', 'sat', 'on', 'the', 'mat', '.', '[PAD]']

visualize_attention(attention_weights, tokens)

6. Полный классификатор внимания

6.1 Классификация текста на основе внимания

Ниже представлена полная модель классификации текста: она использует кодировщик Transformer для извлечения признаков, затем сжимает последовательности переменной длины в векторы фиксированной длины с помощью attention pooling и, наконец, подаёт их в классификатор.

Пример

import torch
import torch.nn as nn
import math

class AttentionClassifier(nn.Module):
"""Модель классификации текста на основе Transformer + attention pooling"""
def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes,
n_heads=4, num_layers=2, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim

# Векторное представление слов
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.pos_encoding = PositionalEncoding(embed_dim, dropout=dropout)

# Энкодер Transformer
encoder_layer = nn.TransformerEncoderLayer(
d_model=embed_dim,
nhead=n_heads,
dim_feedforward=hidden_dim,
dropout=dropout,
batch_first=True,
activation='gelu'
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)

# Пулинг внимания: изучение весов важности для каждой позиции
self.attn_pool = nn.Sequential(
nn.Linear(embed_dim, 1) # Выводит скалярный вес для каждой позиции
)

# Головка классификации
self.classifier = nn.Sequential(
nn.Linear(embed_dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, num_classes)
)

def forward(self, x):
"""
x: [batch, seq_len] — идентификаторы токенов
Возвращает: logits [batch, num_classes], attn_weights [batch, seq_len]
"""
# Встраивание слов + позиционное кодирование
mask = (x == 0) # маска паддинга
x = self.embedding(x) * math.sqrt(self.embed_dim)
x = self.pos_encoding(x)

# Кодирование трансформером
x = self.transformer(x, src_key_padding_mask=mask)

# Пулинг внимания: вычисление весов для каждой позиции, взвешенная сумма
scores = self.attn_pool(x).squeeze(-1) # [batch, seq_len]
scores = scores.masked_fill(mask, float('-inf')) # Игнорировать паддинг
attn_weights = torch.softmax(scores, dim=1) # [batch, seq_len]
pooled = torch.bmm(attn_weights.unsqueeze(1), x).squeeze(1) # [batch, embed_dim]

# Классификация
logits = self.classifier(pooled)
return logits, attn_weights

#
model = AttentionClassifier(
vocab_size=10000, embed_dim=128, hidden_dim=256, num_classes=5
)

# Имитированный вход (с паддингом)
x = torch.randint(1, 10000, (16, 50))
x[:, -10:] = 0 # Последние 10 позиций — это паддинг

logits, attn_weights = model(x)
print(f"Ввод: {x.shape}")
print(f"Логиты: {logits.shape}") # [16, 5]
print(f"Веса внимания: {attn_weights.shape}") # [16, 50]

6.2 Сравнение распространённых вариантов внимания

Type	Источник Q	Источник K/V	Основное использование	Типичное применение
Самовнимание	Сам по себе	Сам по себе	Глобальное взаимодействие внутри последовательности	Transformer、BERT、ViT
Перекрёстное внимание	Целевая последовательность	Исходная последовательность	Слияние информации между последовательностями	Машинный перевод, генерация подписей к изображениям, Stable Diffusion
Канальное внимание	Признаки глобального пулинга	То же, что выше	Изучение важности каналов	SENet、CBAM
Пространственное внимание	Признаки, сжатые по каналам	То же, что выше	Изучение важности пространственных позиций	CBAM, Spatial Transformer
Разреженное внимание	Частичные позиции	Частичные позиции	Снижение вычислительной сложности для длинных последовательностей	Longformer, BigBird

7. Лучшие практики и типичные проблемы

7.1 Советы по использованию

Согласованность масок: Убедитесь, что маска padding и causal-маска имеют согласованное направление (True = игнорировать vs 1 = сохранять). В разных API соглашения могут различаться.
Выбор количества голов: Обычно задаётся как делитель $$d_{model}$$, часто 8 или 16. Слишком много голов уменьшает размерность на одну голову $$d_k$$, что может ухудшить производительность.
Остаточная связь: Обязательно используйте остаточные связи + нормализацию слоя до и после слоёв внимания; это необходимое условие для обучения глубоких трансформеров.
Разогрев скорости обучения: На ранней стадии обучения трансформера колебания градиента велики; использование планировщика с разогревом + косинусным затуханием может значительно повысить стабильность.
Flash Attention: PyTorch 2.0+ поддерживает черезnn.functional.scaled_dot_product_attentionавтоматический вызов Flash Attention, значительно оптимизируя как использование памяти, так и скорость.

7.2 Типичные проблемы

Проблема	Причина	Решение
Все веса внимания равномерно распределены	Ненормализованный вход или слишком большая скорость обучения	Проверьте масштаб входных данных, используйте LayerNorm, уменьшите скорость обучения
Потери при обучении не уменьшаются	Отсутствует позиционное кодирование или неверная маска	Убедитесь, что позиционное кодирование добавлено, проверьте направление маски
Нехватка памяти (длинные последовательности)	Матрица внимания имеет сложность $$O(n^2)$$	Используйте Flash Attention, градиентный чекпоинтинг или разреженное внимание
Деградация внимания при инференсе	Несогласованная маска между обучением и инференсом	Убедитесь, что маска паддинга использует ту же логику при обучении и инференсе

Механизм внимания — ключевая технология современного глубокого обучения. От NLP до CV, от речи до предсказания структуры белков — почти все передовые модели основаны на внимании. Понимание его принципов и реализации — первый ключевой шаг к освоению таких моделей, как Transformer, ViT, Stable Diffusion и GPT.

Другие расширения