Реализация MultiHeadAttention

- Добавлен класс MultiHeadAttention
- Создана документация с блок-схемой
- Добавлен пример использования
- Обновлен README.md

simple_llm/transformer/multi_head_attention.py

@@ -0,0 +1,104 @@
+from torch import nn
+import torch
+from simple_llm.transformer.head_attention import HeadAttention
+
+class MultiHeadAttention(nn.Module):
+    """
+    Реализация механизма многоголового внимания (Multi-Head Attention) из архитектуры Transformer.
+
+    Основные характеристики:
+    - Параллельная обработка входных данных несколькими головами внимания
+    - Поддержка маскирования (causal mask и пользовательские маски)
+    - Финальная проекция с dropout регуляризацией
+
+    Математическое описание:
+    MultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^O
+    где head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)
+
+    Примеры использования:
+
+    1. Базовый пример:
+    >>> mha = MultiHeadAttention(num_heads=8, emb_size=512, head_size=64, max_seq_len=1024)
+    >>> x = torch.randn(2, 50, 512)  # [batch_size, seq_len, emb_size]
+    >>> output = mha(x)  # [2, 50, 512]
+
+    2. С использованием маски:
+    >>> mask = torch.tril(torch.ones(50, 50))  # Causal mask
+    >>> output = mha(x, mask)
+
+    3. Интеграция в Transformer:
+    >>> # В составе Transformer слоя
+    >>> self.attention = MultiHeadAttention(...)
+    >>> x = self.attention(x, mask)
+    """
+    def __init__(self, num_heads: int, emb_size: int, head_size: int, max_seq_len: int, dropout: float = 0.1):
+        """
+        Инициализация многоголового внимания.
+
+        Параметры:
+            num_heads (int): Количество голов внимания. Типичные значения: 4-16
+            emb_size (int): Размерность входных и выходных эмбеддингов
+            head_size (int): Размерность каждой головы внимания (обычно emb_size // num_heads)
+            max_seq_len (int): Максимальная длина последовательности
+            dropout (float): Вероятность dropout (по умолчанию 0.1)
+
+        Контрольные значения:
+            - num_heads * head_size должно равняться emb_size
+            - head_size обычно выбирают 32-128
+            - max_seq_len зависит от задачи (512 для BERT, 2048 для GPT-3)
+        """
+        super().__init__()
+        self._heads = nn.ModuleList([
+            HeadAttention(
+                emb_size=emb_size, 
+                head_size=head_size, 
+                max_seq_len=max_seq_len
+            ) for _ in range(num_heads)
+        ])
+        self._layer = nn.Linear(head_size * num_heads, emb_size)
+        self._dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
+        """
+        Прямой проход через слой многоголового внимания.
+
+        Подробное описание преобразований тензоров:
+        1. Входной тензор [batch_size, seq_len, emb_size] разделяется на N голов:
+           - Каждая голова получает тензор [batch_size, seq_len, head_size]
+        2. Каждая голова вычисляет attention:
+           - Вход: [batch_size, seq_len, head_size]
+           - Выход: [batch_size, seq_len, head_size]
+        3. Конкатенация результатов:
+           - Объединенный выход: [batch_size, seq_len, num_heads * head_size]
+        4. Линейная проекция:
+           - Выход: [batch_size, seq_len, emb_size]
+        5. Применение dropout
+
+        Аргументы:
+            x (torch.Tensor): Входной тензор формы [batch_size, seq_len, emb_size]
+            mask (torch.Tensor, optional): Маска внимания формы [seq_len, seq_len]
+
+        Возвращает:
+            torch.Tensor: Выходной тензор формы [batch_size, seq_len, emb_size]
+
+        Пример преобразований для emb_size=512, num_heads=8:
+        Вход: [4, 100, 512]
+        -> Каждая голова: [4, 100, 64]
+        -> После внимания: 8 x [4, 100, 64] 
+        -> Конкатенация: [4, 100, 512]
+        -> Проекция: [4, 100, 512]
+        -> Dropout: [4, 100, 512]
+        """
+        # 1. Вычисляем attention для каждой головы
+        attention_outputs = [head(x) for head in self._heads]
+        
+        # 2. Объединяем результаты всех голов
+        concatenated_attention = torch.cat(attention_outputs, dim=-1)
+        
+        # 3. Проецируем в пространство эмбеддингов
+        projected_output = self._layer(concatenated_attention)
+        
+        # 4. Применяем dropout для регуляризации
+        final_output = self._dropout(projected_output)
+        
+        return final_output