docs(core): improve and expand docstrings for Decoder module

- docs: rewrite and expand docstrings for Decoder class and its methods (__init__, forward)
- clarify the block’s architecture, pre-LN logic, flow with residual connections, and attention masking
- add mathematical pseudocode, motivation, feature list, usage example, and external references (papers, blog)
- no logic or behavior changes

This improves readability and makes the codebase easier to understand for transformer/LLM practitioners.
This commit is contained in:
Sergey Penkovsky
2025-10-16 12:40:46 +03:00
parent 923aa51e2a
commit 3a356f5d79

View File

@@ -6,37 +6,45 @@ from .multi_head_attention import MultiHeadAttention
class Decoder(nn.Module):
"""
Базовый автогерессивный блок-декодер трансформера (без кэша KV).
Decoder — базовый transformer decoder block (pre-LN), классический строительный блок современных языковых моделей.
Предназначен для:
- Обработки последовательностей с учетом контекста (самовнимание)
- Постепенного генерирования выходной последовательности
- Учета масок для предотвращения "заглядывания в будущее"
Назначение:
-----------
- Инкапсулирует архитектуру: norm → multi-head self-attention → residual → norm → feed-forward → residual
- Подходит как для LLM/GPT, так и для любых autoregressive sequence моделей.
- Использует masked self-attention: каждый токен видит только предыдущие (никакого \"заглядывания в будущее\").
- Стабильность обеспечивается через residual connections и LayerNorm после каждого sub-layer.
Алгоритм работы:
1. Входной тензор (batch_size, seq_len, emb_size)
2. Многоголовое внимание с residual connection и LayerNorm
3. FeedForward сеть с residual connection и LayerNorm
4. Выходной тензор (batch_size, seq_len, emb_size)
Почему это важно?
-----------------
- Все современные языковые модели состоят из подобных блоков, соединённых в стек.
- Алгоритм residual+norm позволяет проще обучать очень глубокие сети.
- Разделение на attention+FFN дает и локальные, и глобальные взаимодействия между токенами.
Основные характеристики:
- Поддержка масок внимания
- Residual connections для стабилизации градиентов
- Layer Normalization после каждого sub-layer
- Конфигурируемые параметры внимания
Формула работы (псевдокод):
---------------------------
y1 = norm1(x)
attn_out = Attention(y1)
x2 = x + attn_out # residual
y2 = norm2(x2)
ffn_out = FFN(y2)
out = x2 + ffn_out # residual
Архитектурные особенности:
--------------------------
- Поддержка внимания с маской (causal mask или произвольная attention mask)
- Residual connections для каждого блока (attention, FFN)
- Pre-LN (norm перед каждым подблоком)
- Зависит от переданных блоков self_attention и feed_forward, а не их реализации
References:
-----------
- Vaswani et al., \"Attention is All You Need\" (2017): https://arxiv.org/abs/1706.03762
- Illustrated Transformer: https://jalammar.github.io/illustrated-transformer/
- Transformer Circuits (дружественное описание): https://transformer-circuits.pub/2021/framework/index.html
Научная суть:
- Осуществляет посимвольное предсказание: каждый токен видит только предыдущие (masked attention)
- Состоит из self-attention + feedforward + residual + нормализация
- Residual connection и normalization дают стабильность и градиентный “flow” при обучении
- Механизм предложен в Vaswani et al., "Attention is All You Need", 2017
Args:
num_heads (int): количество attention-голов
emb_size (int): размер эмбеддинга
head_size (int): размер одной attention-головы
max_seq_len (int): максимальная длина последовательности
dropout (float): вероятность dropout
Пример:
-------
>>> decoder = Decoder(num_heads=8, emb_size=512, head_size=64, max_seq_len=1024)
>>> x = torch.randn(1, 10, 512)
>>> out = decoder(x)
@@ -52,14 +60,27 @@ class Decoder(nn.Module):
dropout: float = 0.1,
):
"""
Инициализация декодера.
Инициализация стандартного decoder-блока для Transformer.
Параметры:
num_heads: int - количество голов внимания
emb_size: int - размерность эмбеддингов
head_size: int - размерность каждой головы внимания
max_seq_len: int - максимальная длина последовательности
dropout: float (default=0.1) - вероятность dropout
Аргументы:
----------
num_heads: int
Количество attention голов (как делить emb_size на heads)
emb_size: int
Размерность эмбеддингов (и входа и выхода)
head_size: int
Размерность одной attention-головы (emb_size = num_heads * head_size)
max_seq_len: int
Максимальная длина последовательности (важно для mask)
dropout: float, default=0.1
Dropout после внимания и FFN
Внутри:
-------
- Создаёт слой MultiHeadAttention (masked/casual)
- Создаёт двухслойный FeedForward (SwiGLU или GELU)
- Применяет 2 слоя LayerNorm для стабилизации градиентов
- Все блоки реализованы как PyTorch-модули
"""
super().__init__()
self._heads = MultiHeadAttention(
@@ -75,20 +96,26 @@ class Decoder(nn.Module):
def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
"""
Прямой проход через декодер.
Один прямой проход через Transformer decoder block.
Вход:
x: torch.Tensor - входной тензор [batch_size, seq_len, emb_size]
mask: torch.Tensor (optional) - маска внимания [seq_len, seq_len]
Аргументы:
----------
x : torch.Tensor
Входной тензор [batch_size, seq_len, emb_size]
mask : torch.Tensor, optional
Attention/causal mask (по умолчанию None, тогда будет casual mask по длине seq_len)
Возвращает:
torch.Tensor - выходной тензор [batch_size, seq_len, emb_size]
-----------
out : torch.Tensor
Выходной тензор той же формы, что и x
Алгоритм forward:
1. Применяем MultiHeadAttention к входу
2. Добавляем residual connection и LayerNorm
3. Применяем FeedForward сеть
4. Добавляем residual connection и LayerNorm
Алгоритм:
---------
- Применяем attention к нормализованному входу (layernorm)
- Добавляем residual-связь (attention + исходный вход)
- Применяем FFN к нормализованному результату (layernorm)
- Добавляем residual-связь (ffn + предыдущий выход)
"""
# Self-Attention блок
attention, _ = self._heads(x, mask, use_cache=False, cache=None)