AI Research Foundations, курс 4, часть 2: позиции и параметры трансформера

Перевод лабораторных 4.4–4.5 курса Google DeepMind «Discover the Transformer Architecture»: почему внимание без позиционных эмбеддингов не видит порядка слов и сколько обучаемых параметров в каждом блоке трансформера.

Вторая половина курса 4 «Discover the Transformer Architecture». В первой части вы посчитали внимание по матрицам Gemma-1B. Здесь два вопроса, без которых трансформер не понять до конца: откуда модель знает порядок слов, если внимание его не видит, и из чего складываются миллионы (и миллиарды) параметров.

Как проходить вместе с курсом ChinaAI

Лабораторная	Время	Когда делать	Что нужно
4.4. Позиционные эмбеддинги	15 мин	Перед уроком «Позиционное кодирование и RoPE»	NumPy, SciPy, CPU
4.5. Обучаемые параметры трансформера	25 мин + 1–10 мин обучения	Вместе с уроками «Блок трансформера» и «GPT на PyTorch»	Keras на JAX; для обучения желательно GPU

Трансформер в лабораторной 4.5 — «классический» (Post-LN, LayerNorm, ReLU, синусоидальные позиции, как в статье 2017 года). Современные модели устроены немного иначе: RMSNorm перед подблоком, SwiGLU, RoPE, общие ключи у голов. После лабораторной прочитайте урок «Современный блок LLM» и пересчитайте параметры для такого блока — отличное упражнение.

Лабораторная 4.4. Позиционные эмбеддинги

Ноутбук: gdm_lab_4_4_positional_embeddings.ipynb · открыть в Colab

Посмотрите, как self-attention без позиционных эмбеддингов не зависит от порядка. 15 минут.

Обзор

В этой лабораторной вы увидите, что механизм внимания инвариантен к позициям (position invariant): выход механизма внимания один и тот же, в каком бы порядке ни шли предыдущие токены.

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

Чему вы научитесь

К концу лабораторной вы поймёте, почему механизм self-attention даёт тот же выход, если переставить токены в промпте.

Задачи

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

Импорты

import numpy as np  # Чтобы задавать эмбеддинги и работать с ними.
from scipy.special import softmax  # Чтобы вычислять веса внимания.

Готовим модель

Запускайте эти ячейки по одной, чтобы вычислить выход головы внимания и соответствующие веса внимания.

Следующий блок кода сначала задаёт предложения и словарь модели. Модель игрушечная, поэтому словарь ограничен пятью токенами из словаря vocabulary (плюс служебный <pad>).

# Задаём предложения. Всё в нижнем регистре, чтобы у "The" и "the" был один
# и тот же эмбеддинг.
sentence1_str = "the zebra chased the lion ."   # «зебра гналась за львом»
sentence2_str = "the lion chased the zebra ."   # «лев гнался за зеброй»

# Задаём словарь.
vocabulary = {"<pad>": 0, "the": 1, "zebra": 2, "chased": 3, "lion": 4, ".": 5}
inv_vocabulary = {v: k for k, v in vocabulary.items()}

Следующий блок строит случайную матрицу эмбеддингов для этих токенов. Такая матрица не отражает никакого смыслового сходства между словами. Цель упражнения — показать инвариантность внимания к позициям, а не использовать эти эмбеддинги в модели для предсказания следующего токена. Поэтому неважно, что для предсказаний эти эмбеддинги не годятся.

# Размерность эмбеддингов.
embedding_dim = 3
vocabulary_size = len(vocabulary)

# Фиксируем seed для воспроизводимости.
np.random.seed(2311)

# Матрица эмбеддингов (vocab_size x d_model).
embedding_matrix = np.random.rand(vocabulary_size, embedding_dim)

Следующий блок так же строит случайные матрицы проекций запросов, ключей и значений. В настоящей модели эти матрицы выучиваются по данным. Для этого упражнения опять же неважно, что параметры не несут ничего полезного для предсказаний.

# Размерность векторов ключей, запросов и значений (может отличаться от
# embedding_dim, но здесь для простоты та же).
d_k = embedding_dim
d_q = embedding_dim
d_v = embedding_dim

# Константа для замаскированных оценок внимания.
K_MASK = -2.3819763e38  # Очень большое по модулю отрицательное число.

# Матрицы проекций (embedding_dim x d_k для W_q, embedding_dim x d_q для W_k,
# embedding_dim x d_v для W_v).
W_q = np.random.rand(embedding_dim, d_k)
W_k = np.random.rand(embedding_dim, d_q)
W_v = np.random.rand(embedding_dim, d_v)

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

def compute_attention_output(sentence: str) -> tuple[np.ndarray, np.ndarray]:
    """Вычисляет выход внимания и веса внимания для всех токенов `sentence`.

    Аргументы:
      sentence: предложение, для которого считаются веса внимания. Токены
        разделены пробелами и должны быть из словаря vocabulary.

    Возвращает:
      attention_output: выход механизма внимания.
        Форма: (num_tokens, d_v).
      attention_weights: веса внимания для всех токенов.
        Форма: (num_tokens, num_tokens).
    """

    tokens = [vocabulary[word] for word in sentence.split()]

    embeddings = embedding_matrix[tokens]  # обратите внимание: никакой информации о позиции

    # Вычисляем запросы, ключи и значения для предложения.
    Q = embeddings @ W_q  # Форма: (num_tokens, d_q)
    K = embeddings @ W_k  # Форма: (num_tokens, d_k)
    V = embeddings @ W_v  # Форма: (num_tokens, d_v)

    # Вычисляем маску внимания.
    l = len(tokens)
    attention_mask = np.tri(l)  # Форма: (num_tokens, num_tokens).
    print(f"\n--- Sentence: \"{sentence}\" ---")

    # Вычисляем логиты внимания с нормирующим множителем.
    scale_factor = np.sqrt(d_k)
    logits = (Q @ K.T) / scale_factor  # Форма: (num_tokens, num_tokens).
    # Применяем маску внимания.
    # Форма: (num_tokens, num_tokens).
    logits = np.where(attention_mask, logits, K_MASK)

    print("\nAttention logits (Q @ K.T / sqrt(d_k)) for last token:\n",
          logits[-1, :])

    # Вычисляем веса внимания (SoftMax).
    # Форма: (num_tokens, num_tokens).
    attention_weights = softmax(logits, axis=1)

    print("\nAttention weights (SoftMax) for last token:\n",
          attention_weights[-1, :])

    print("\nValue matrix:\n", V)

    # Вычисляем выход внимания.
    attention_output = attention_weights @ V  # Форма: (num_tokens, d_v)
    print("\nAttention output (weights @ V) for last token:\n",
          attention_output[-1, :])

    return attention_output, attention_weights

Сравниваем предложения с разным порядком слов

Теперь можно вычислить веса и выходы внимания для предложений «the zebra chased the lion .» и «the lion chased the zebra .».

attention_output1, attention_weights1 = compute_attention_output(sentence1_str)
attention_output2, attention_weights2 = compute_attention_output(sentence2_str)

Что вы заметили?

Посмотрите внимательно на веса внимания и матрицы значений выше.

Как вы увидите, и веса внимания, и матрицы значений у двух предложений различаются. В векторе весов внимания второй и пятый элементы поменялись местами (они соответствуют «zebra» и «lion»).

Две матрицы значений различаются во второй и пятой строках — это представления «zebra» и «lion».

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

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

Ожидаемый вывод (наш запуск, округлено): веса внимания последнего токена — [0.1388 0.1412 0.1058 0.1388 0.302 0.1734] для первого предложения и [0.1388 0.302 0.1058 0.1388 0.1412 0.1734] для второго; выход для последнего токена в обоих случаях [1.1482 1.2014 1.0222].

Почему так? Выход — это сумма $\sum_j \alpha_j v_j$ по всем токенам. Пара «вес токена, вектор значения токена» зависит только от самого токена (и от последнего токена-запроса), а не от того, где он стоит. От перестановки слагаемых сумма не меняется: $0{,}14 \cdot v_{\text{zebra}} + 0{,}30 \cdot v_{\text{lion}}$ в первом предложении и $0{,}30 \cdot v_{\text{lion}} + 0{,}14 \cdot v_{\text{zebra}}$ во втором — одно и то же.

Запустите следующую ячейку, чтобы увидеть это сравнение ещё нагляднее.

print("\n--- Comparison ---")

print("\nAttention weights for last token in S1:", attention_weights1[-1, :])
print("Attention weights for last token in S2:", attention_weights2[-1, :])

print("\nAttention outputs for last token in S1:", attention_output1[-1, :])
print("Attention outputs for last token in S2:", attention_output2[-1, :])


print("\nAre attention outputs for last token in S1 and S2 the same?",
      np.allclose(attention_output1[-1, :], attention_output2[-1, :]))

Сравниваем веса внимания для других предложений

Чтобы убедиться, что это не случайность двух предложений выше, сравните и другие пары предложений.

# @title Сравниваем веса и выходы внимания для других предложений

sentence_1 = "lion the chased zebra the ."  # @param {"type": "string"}
sentence_2 = "chased lion the the zebra ."  # @param {"type": "string"}

tokens_1 = sentence_1.split()
tokens_2 = sentence_2.split()

possible_token_list = "', '".join(vocabulary.keys())

# Проверяем, что все токены есть в словаре.
for t in tokens_1:
    if t not in vocabulary:
        raise ValueError(
            f"Invalid token '{t}' in sentence_1. Please only use one of the"
            f" following tokens: '{possible_token_list}'."
        )

for t in tokens_2:
    if t not in vocabulary:
        raise ValueError(
            f"Invalid token '{t}' in sentence_2. Please only use one of the"
            f" following tokens: '{possible_token_list}'."
        )

attention_output1, attention_weights1 = compute_attention_output(sentence_1)
attention_output2, attention_weights2 = compute_attention_output(sentence_2)

print("\n--- Comparison ---")

print("\nAttention weights for last token in S1:", attention_weights1[-1, :])
print("Attention weights for last token in S2:", attention_weights2[-1, :])

print("\nAttention outputs for last token in S1:", attention_output1[-1, :])
print("Attention outputs for last token in S2:", attention_output2[-1, :])

print(
    "\nAre attention outputs for last token in S1 and S2 the same?",
    np.allclose(attention_output1[-1, :], attention_output2[-1, :]),
)

Для пары по умолчанию (последний токен у обоих — «.») выходы совпадают: True. А для «lion the chased .» и «chased the . lion» у нас получилось False: выходы [1.2899 1.3412 1.1294] и [1.3592 1.422 1.2099]. Причина — последний токен задаёт запрос: веса внимания — это скалярные произведения запроса последнего токена с ключами. Другой последний токен — другой запрос, другие веса и другой выход. Инвариантность касается только порядка предыдущих токенов.

Итоги

Это интерактивное задание показало, что выход механизма внимания для последнего токена (то есть эмбеддинг, по которому модель предсказывает следующий токен) не зависит от порядка токенов в промпте. Значит, механизм внимания инвариантен к порядку (order-invariant).

В следующем задании вы узнаете о способах кодировать информацию о позиции в трансформерах.

Это добавление переводчика. Чтобы внимание различало «зебра гналась за львом» и «лев гнался за зеброй», к эмбеддингам добавляют информацию о позиции. В лабораторной 4.5 это синусоидальные позиционные эмбеддинги, которые просто прибавляются к эмбеддингам токенов. Современные модели (Gemma, Qwen, Llama) используют RoPE — поворот векторов запросов и ключей на угол, зависящий от позиции; см. урок «Позиционное кодирование и RoPE». Проверьте сами: прибавьте в compute_attention_output к embeddings слагаемое np.arange(len(tokens))[:, None] * 0.1 — и выходы двух предложений станут разными.

Лабораторная 4.5. Обучаемые параметры трансформера

Ноутбук: gdm_lab_4_5_reflection_on_trainable_parameters.ipynb · открыть в Colab

Посмотрите, как отдельные компоненты нейросети собираются в полный трансформер. 25 минут.

Обзор

В этой лабораторной вы пройдёте по полной реализации трансформера на Keras и JAX, компонент за компонентом. По ходу вы напишете функции, которые считают число обучаемых параметров каждого компонента, а в конце — число обучаемых параметров всего трансформера.

Лабораторная даёт возможность разобрать все детали полной архитектуры трансформера. Кроме того, она показывает важное свойство реализаций сложных нейросетей — модульность (modularity). Если разбить сложную архитектуру на небольшие кубики и затем собрать их вместе, с ней гораздо проще управляться, и отдельные компоненты легко менять.

Чему вы научитесь

К концу лабораторной вы:

поймёте, из каких компонентов нейросети собирается трансформер;
поймёте, сколько обучаемых параметров у каждого из этих компонентов;
поймёте, как гиперпараметры вроде числа блоков трансформера или размера словаря влияют на общее число параметров модели.

Задачи

В этой лабораторной вы:

разберёте компоненты нейросети, из которых собирается трансформер:
1. нормализация слоя (layer normalization);
2. слой эмбеддингов;
3. многоголовое внимание;
4. многослойный перцептрон;
5. блок трансформера;
6. выходной слой;
7. полный трансформер;
для каждого компонента реализуете функцию, которая считает число его обучаемых параметров, — чтобы глубже понять компонент;
поменяете гиперпараметры модели и посмотрите, как размер разных компонентов влияет на общее число параметров;
(необязательно) обучите реализованную в лабораторной модель на наборе Africa Galore и убедитесь, что реализация работает как надо.

flowchart TB
IN["id токенов"] --> EMB["TokenAndPositionEmbedding<br/>V × E параметров"]
EMB --> B1["TransformerBlock 1"]
B1 --> B2["TransformerBlock 2"]
B2 --> OUT["OutputLayer<br/>(E + 1) × V параметров"]
subgraph BLK["внутри блока"]
  A["MultiHeadSelfAttention<br/>4 × (E + 1) × E"] --> R1["+ остаточная связь, LayerNorm (2E)"]
  R1 --> M["MultiLayerPerceptron<br/>(E + 1)·M + (M + 1)·E"]
  M --> R2["+ остаточная связь, LayerNorm (2E)"]
end

Устройство трансформера лабораторной 4.5 и формулы числа параметров (E — embedding_dim, M — mlp_dim, V — размер словаря)

Импорты

В этой лабораторной вы в основном пишете функции с простыми вычислениями, которым не нужны дополнительные пакеты. Уже реализованный в лабораторной трансформер использует пакеты Keras и JAX, а методы пакета курса ai_foundations проверяют вашу реализацию.

%%capture
!pip install jax[cuda12]==0.7.2
# Ставим специальный пакет курса.
!pip install "git+https://github.com/google-deepmind/ai-foundations.git@main"

import os  # Чтобы задать параметры Keras.
os.environ["KERAS_BACKEND"] = "jax"

import jax  # Для работы с векторами и матрицами.
import jax.numpy as jnp  # Для работы с векторами и матрицами.
import pandas as pd  # Чтобы загрузить набор данных.
import keras  # Чтобы задать трансформер.
import tqdm  # Чтобы показывать индикаторы прогресса.
from keras import layers  # Чтобы задать трансформер.
from ai_foundations import training  # Чтобы обучить трансформер.
from ai_foundations import generation  # Чтобы давать промпты трансформеру.
from ai_foundations import tokenization  # Чтобы загрузить BPE-токенизатор.
# Для обратной связи.
from ai_foundations.feedback.course_4 import counting_parameters as feedback

Компоненты трансформера

В следующих ячейках на Keras реализованы все компоненты трансформера. Пройдите по каждому компоненту, а в конце каждого реализуйте функцию, которая считает, сколько параметров обучается в этом компоненте.

Нормализация слоя

Ячейка ниже реализует нормализацию слоя на Keras и JAX. Напомним формулу нормализации слоя:

\text{LayerNorm}(\mathbf{x}) = \gamma \frac{\mathbf{x} - \mu(\mathbf{x})}{\sqrt{\text{Var}(\mathbf{x}) + \epsilon}} + \beta,

где $\mu(\mathbf{x})$ — среднее элементов $\mathbf{x}$ , а $\text{Var}(\mathbf{x})$ — дисперсия элементов $\mathbf{x}$ . $\gamma$ (гамма) и $\beta$ (бета) — обучаемые параметры.

Словами: из каждого вектора вычитаем его среднее и делим на его стандартное отклонение — получаем вектор со средним 0 и разбросом 1, — а затем растягиваем его на $\gamma$ и сдвигаем на $\beta$ . Например, для $\mathbf{x} = (1, 2, 3)$ : среднее 2, дисперсия $2/3$ , нормированный вектор примерно $(-1{,}22;\ 0;\ 1{,}22)$ .

Чтобы понять, как отдельные компоненты нейросети реализуются в Keras, заметьте, что у всех слоёв Keras есть методы __init__ и call.

Метод инициализации (__init__). Он вызывается с настройками компонента. Например, у нормализации слоя две настройки: размер эмбеддингов embedding_dim и аргумент epsilon — константа, которая прибавляется к дисперсии, чтобы не делить на 0.

По аргументу embedding_dim метод создаёт векторы $\gamma$ и $\beta$ — основные параметры этого компонента.

Метод call. Метод call (его ещё называют прямым проходом, forward) вызывается каждый раз, когда модель делает предсказание: при обучении — чтобы вычислить функцию потерь, при валидации и тестировании — чтобы предсказать на точках, которых нет в обучающих данных. Он определяет, как вход превращается в выход компонента с помощью параметров модели.

Посмотрите на функцию call ниже. Она вызывается со входом компонента x и возвращает нормализованные входы.

Функция выполняет такие шаги:

Вычисляет среднее $\mu(\mathbf{x})$ по всем признакам функцией jnp.mean. Обратите внимание: аргумент axis равен -1. Это говорит jnp.mean, что среднее нужно брать по последнему измерению x, где по соглашению лежат все признаки одной точки данных.
Вычисляет дисперсию по всем признакам по формуле $\text{Var}(\mathbf{x}) = \frac{1}{d_x} \sum \left(\mathbf{x} - \mu(\mathbf{x})\right)^2$ , где $d_x$ — размерность x.
Вычисляет нормализованные значения x по среднему и дисперсии.
Возвращает нормализованные значения, умноженные на масштаб $\gamma$ и сдвинутые на $\beta$ .

Запустите следующую ячейку, чтобы задать компонент LayerNorm. Затем пройдите по коду строка за строкой.

class LayerNorm(layers.Layer):
    """Реализация нормализации слоя (Layer Normalization) на Keras.

    Нормализует активации предыдущего слоя для каждого примера батча
    отдельно, по измерению признаков.

    Аргументы:
      embedding_dim: размерность выхода механизма внимания.
      epsilon: маленькое число, которое прибавляется к дисперсии,
        чтобы не делить на ноль.
    """

    def __init__(self, embedding_dim: int, epsilon: float = 1e-6, **kwargs):
        super().__init__(**kwargs)
        self.epsilon = epsilon

        # Форма параметров (gamma и beta) — это размер входных эмбеддингов
        # и выхода механизма внимания `embedding_dim`.
        shape = (embedding_dim,)

        # Инициализируем gamma (масштаб) вектором из единиц.
        self.gamma = self.add_weight(
            shape=shape,
            initializer='ones',
            name='gamma'
        )
        # Инициализируем beta (сдвиг) вектором из нулей.
        self.beta = self.add_weight(
            shape=shape,
            initializer='zeros',
            name='beta'
        )

    def call(self, x: jax.Array):
        """
        Применяет нормализацию слоя.

        Аргументы:
          x: входной тензор.
            Форма: (batch_size, sequence_length, embedding_dim).

        Возвращает:
          Нормализованный и преобразованный тензор.
            Форма: (batch_size, sequence_length, embedding_dim).
        """
        # Считаем среднее и дисперсию по оси признаков (-1).
        mean = jnp.mean(x, axis=-1, keepdims=True)
        variance = jnp.mean(jnp.square(x - mean), axis=-1, keepdims=True)

        # Нормализуем вход.
        normalized_x = (x - mean) / jnp.sqrt(variance + self.epsilon)

        # Применяем выученные масштаб (gamma) и сдвиг (beta).
        return self.gamma * normalized_x + self.beta

Сколько обучаемых параметров у этого компонента? Здесь и во всей лабораторной считайте, что модель создана с такими гиперпараметрами:

{
    # Максимальное число входных токенов в последовательности.
    "max_length": 128,
    # Размерность входных эмбеддингов, выходов механизма внимания
    # и выхода каждого блока трансформера.
    "embedding_dim": 256,
    # Размерность скрытого слоя MLP в каждом блоке трансформера.
    "mlp_dim": 384,
    # Число голов внимания.
    "num_heads": 4,
    # Число блоков трансформера.
    "num_blocks": 2,
    # Число уникальных токенов в словаре.
    "vocabulary_size": 262144
}

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

# Задаём гиперпараметры модели.
MODEL_HYPERPARAMETERS = {
    "max_length": 128,
    "embedding_dim": 256,
    "mlp_dim": 384,
    "num_heads": 4,
    "num_blocks": 2,
    "vocabulary_size": 262144
}


def parameter_count_layer_norm(hyperparams: dict[str, int]) -> int:
    embedding_dim =  hyperparams["embedding_dim"]
    parameter_count = embedding_dim + embedding_dim   # gamma + beta
    return parameter_count


param_count = parameter_count_layer_norm(MODEL_HYPERPARAMETERS)
if param_count != ...:
    print(f"Number of parameters for layer normalization: {param_count:,}")

# Запустите эту ячейку, чтобы проверить `parameter_count_layer_norm`.
feedback.test_parameter_count_layer_norm(parameter_count_layer_norm)

Ожидаемый вывод: Number of parameters for layer normalization: 512.

Задание 1. Эмбеддинги токенов и позиций

Следующая ячейка реализует и синусоидальные позиционные эмбеддинги, и компонент, который превращает входные токены в эмбеддинги. У синусоидальных позиционных эмбеддингов нет обучаемых параметров, так что функцию _positional_encoding можно не учитывать.

class TokenAndPositionEmbedding(layers.Layer):
    """
    Слой Keras, который объединяет эмбеддинги токенов и синусоидальные
    позиционные эмбеддинги.

    Аргументы:
      vocabulary_size: размер словаря.
      embedding_dim: размерность эмбеддингов. Должна быть чётной.
      max_length: максимальная длина входных последовательностей в токенах.
    """

    def __init__(
        self,
        vocabulary_size: int,
        embedding_dim: int,
        max_length: int,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.vocabulary_size = vocabulary_size
        self.embedding_dim = embedding_dim
        self.max_length = max_length
        self.pos_encoding = self._positional_encoding()  # фиксированная, не обучается

        # Таблица эмбеддингов токенов — вот здесь обучаемые параметры.
        self.token_embeddings = layers.Embedding(
            input_dim=vocabulary_size, output_dim=embedding_dim, mask_zero=True
        )

    def _positional_encoding(self) -> jax.Array:
        """
        Создаёт фиксированную матрицу синусоидального позиционного кодирования.

        Строит уникальное представление позиции каждого токена
        последовательности из синусов и косинусов разных частот.

        Возвращает:
          Массив JAX формы (1, max_length, embedding_dim) с позиционным
            кодированием.
        """
        depth = self.embedding_dim // 2
        # Форма: (max_length, 1).
        positions = jnp.arange(self.max_length)[:, jnp.newaxis]
        depths = jnp.arange(depth)[jnp.newaxis, :] / depth
        # Форма: (1, depth).
        angle_rates = 1 / (10000**depths)
        # Форма: (1, depth).
        angle_rads = positions * angle_rates
        # Форма: (max_length, depth).
        pos_encoding = jnp.concatenate(
            [jnp.sin(angle_rads), jnp.cos(angle_rads)], axis=-1
        )
        # Добавляем измерение батча для broadcasting.
        return pos_encoding[jnp.newaxis, :, :]

    def call(self, x: jax.Array) -> jax.Array:
        """
        Применяет слой эмбеддингов.

        Аргументы:
          x: входной тензор id токенов. Форма: (batch_size, sequence_length).

        Возвращает:
          Массив JAX формы (batch_size, sequence_length, embedding_dim) —
            сумму эмбеддингов токенов и позиционных эмбеддингов.
        """
        # Берём эмбеддинги токенов из таблицы.
        # Входной тензор `x` содержит целые id токенов.
        token_embeddings = self.token_embeddings(x)

        # Масштабируем эмбеддинги токенов, как у Vaswani et al., 2017.
        token_embeddings *= jnp.sqrt(self.embedding_dim)

        # Прибавляем фиксированные позиционные эмбеддинги.
        # Позиционное кодирование формы (1, max_length, embedding_dim)
        # «растягивается» (broadcasting) на измерение батча token_embeddings.
        return token_embeddings + self.pos_encoding

# Реализуйте подсчёт числа параметров слоя эмбеддингов:
def parameter_count_embedding(hyperparams: dict[str, int]) -> int:

    parameter_count = ...

    return parameter_count


param_count = parameter_count_embedding(MODEL_HYPERPARAMETERS)
if param_count != ...:
    print(f"Number of parameters for embedding component: {param_count:,}")

# Запустите эту ячейку, чтобы проверить `parameter_count_embedding`.
feedback.test_parameter_count_embedding(parameter_count_embedding)

Задание 2. Многоголовое внимание

Следующая ячейка реализует полный механизм многоголового внимания.

Пройдите по реализации, а затем допишите функцию parameter_count_attention.

Метод call здесь сложный: он поддерживает сильно распараллеленные вычисления, где все вычисления для всех голов внимания и всех примеров батча выполняются параллельно на четырёхмерных тензорах (матрицах с четырьмя измерениями). Чтобы оценить число обучаемых параметров, достаточно понять начало MultiHeadSelfAttention и метод __init__. Так что ничего страшного, если вы не до конца понимаете, что происходит в методе call.

Подсказки:

В оригинале здесь есть подсказка, что в компонент входит нормализация слоя и её параметры надо прибавить через parameter_count_layer_norm. Это неточность: в классе MultiHeadSelfAttention нормализации нет (она в TransformerBlock), и проверка с эталонным решением ожидает число без LayerNorm. (примечание переводчика)
Запросы, ключи и значения проецируются в размерность embedding_dim / num_heads. Голов num_heads, значит, общая размерность проекций по всем головам — embedding_dim / num_heads * num_heads = embedding_dim. Поэтому для подсчёта параметров можно считать, что голова одна и она проецирует всё в размерность embedding_dim.
Реализация Dense в Keras автоматически добавляет смещения. Поэтому если на вход слоя Dense подаётся размерность $q$ , а на выходе размерность $r$ , он создаёт матрицу параметров размера $q \times r$ и вектор смещений размера $r$ . Всего $(q + 1) \times r$ параметров.

K_MASK = -2.3819763e38


class MultiHeadSelfAttention(layers.Layer):
    """Слой Keras: многоголовое self-attention с причинной маской.

    Аргументы:
      embedding_dim: размерность эмбеддингов. Должна делиться на num_heads.
      num_heads: число голов внимания.
      dropout_rate: доля dropout для выхода внимания.
    """

    def __init__(
        self,
        embedding_dim: int,
        num_heads: int,
        dropout_rate: float = 0.0,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        self.dropout_rate = dropout_rate

        # embedding_dim должна делиться на num_heads.
        assert (
            embedding_dim % num_heads == 0
        ), "embedding_dim must be divisible by num_heads"
        self.head_dim = embedding_dim // num_heads

        # Слои проекций для запроса, ключа, значения и итогового выхода.
        # Каждый Dense: (embedding_dim + 1) * embedding_dim параметров.
        self.q_dense = keras.layers.Dense(embedding_dim, name="q_projection")
        self.k_dense = keras.layers.Dense(embedding_dim, name="k_projection")
        self.v_dense = keras.layers.Dense(embedding_dim, name="v_projection")
        self.output_dense = keras.layers.Dense(embedding_dim,
                                               name="output_projection")

        # Слой dropout (параметров нет).
        self.dropout = layers.Dropout(rate=dropout_rate)

    def _split_heads(self, x: jax.Array) -> jax.Array:
        """Делит последнее измерение на (num_heads, head_dim) и переставляет
           измерения так, чтобы итоговая форма была
           (batch_size, num_heads, sequence_length, head_dim).

        Аргументы:
          x: входной тензор формы (batch_size, sequence_length, embedding_dim).

        Возвращает:
          jax.Array формы (batch_size, num_heads, sequence_length, head_dim).
        """
        batch_size, sequence_length, _ = x.shape
        x = x.reshape(
            batch_size, sequence_length, self.num_heads, self.head_dim
        )
        return x.transpose((0, 2, 1, 3))

    def call(self, x: jax.Array, training: bool = False) -> jax.Array:
        """Прямой проход слоя MultiHeadSelfAttention.

        Аргументы:
          x: входной тензор формы (batch_size, sequence_length, embedding_dim).
          training: работает ли слой в режиме обучения (применять dropout)
            или в режиме предсказания.

        Возвращает:
          Выходной тензор формы (batch_size, sequence_length, embedding_dim).
        """
        batch_size, sequence_length, _ = x.shape

        # Проецируем входы в Q, K, V.
        query = self.q_dense(x)  # (batch_size, sequence_length, embedding_dim).
        key = self.k_dense(x)  # (batch_size, sequence_length, embedding_dim).
        value = self.v_dense(x)  # (batch_size, sequence_length, embedding_dim).

        # Меняем форму для многоголового внимания.

        # (batch, num_heads, sequence_length, head_dim).
        query = self._split_heads(query)

        # (batch, num_heads, sequence_length, head_dim).
        key = self._split_heads(key)

        # (batch, num_heads, sequence_length, head_dim).
        value = self._split_heads(value)

        # Масштабированное скалярное внимание.
        # Матричное умножение:
        # (..., sequence_length, d_k) x (..., d_k, sequence_length)
        #    -> (..., sequence_length, sequence_length).
        logits_raw = jnp.matmul(query, key.transpose((0, 1, 3, 2)))
        logits_raw /= jnp.sqrt(self.head_dim)

        # Применяем причинную маску (запрет смотреть вперёд):
        # внимание только к предыдущим позициям.
        causal_mask = jnp.tril(jnp.ones((sequence_length, sequence_length),
                                        dtype=bool))
        logits_masked = jnp.where(causal_mask, logits_raw, K_MASK)

        # SoftMax даёт веса внимания.
        # (batch, num_heads, sequence_length, sequence_length).
        attention_weights = jax.nn.softmax(logits_masked, axis=-1)

        # Применяем веса внимания к значениям.
        # (batch, num_heads, sequence_length, head_dim).
        attention_output = jnp.matmul(attention_weights, value)

        # Склеиваем головы и применяем итоговую проекцию.
        # Обратно к форме (batch_size, sequence_length, num_heads, head_dim).
        attention_output = attention_output.transpose((0, 2, 1, 3))

        # Форма (batch_size, sequence_length, embedding_dim).
        attention_output = attention_output.reshape(
            batch_size, sequence_length, self.embedding_dim
        )
        attention_output = self.output_dense(attention_output)

        # Применяем dropout.
        attention_output = self.dropout(attention_output, training=training)

        return attention_output

# Реализуйте здесь подсчёт числа параметров многоголового внимания:
def parameter_count_attention(hyperparams: dict[str, int]) -> int:

    parameter_count = ...

    return parameter_count


param_count = parameter_count_attention(MODEL_HYPERPARAMETERS)
if param_count != ...:
    print("Number of parameters for multi-head attention component:"
          f" {param_count:,}")

# Запустите эту ячейку, чтобы проверить `parameter_count_attention`.
feedback.test_parameter_count_attention(parameter_count_attention)

Задание 3. Компонент MLP

Следующая ячейка реализует компонент MLP блока трансформера.

class MultiLayerPerceptron(layers.Layer):
    """Компонент «многослойный перцептрон».

    Двухслойный MLP. Вносит нелинейность и помогает модели выучивать
    сложные закономерности.

    Аргументы:
      embedding_dim: размерность пространства эмбеддингов.
      mlp_dim: размерность скрытого слоя MLP (часто больше embedding_dim).
      dropout_rate: доля dropout для выхода MLP.
      activation: функция активации первого полносвязного слоя.

    Возвращает:
      Выходной тензор формы (batch_size, sequence_length, embedding_dim).
    """

    def __init__(self,
                 embedding_dim: int,
                 mlp_dim: int,
                 dropout_rate: float = 0.0,
                 activation: str = "relu",
                 **kwargs: dict):
        super().__init__(**kwargs)
        # Задаём двухслойный MLP.
        self.mlp = keras.Sequential([
            # Скрытый слой: embedding_dim -> mlp_dim.
            layers.Dense(mlp_dim, activation=activation),
            # Выходной слой: mlp_dim -> embedding_dim.
            layers.Dense(embedding_dim)
        ])
        self.dropout = layers.Dropout(dropout_rate)

    def call(self, x: jax.Array) -> jax.Array:
        """Применяет MLP к входному тензору.

        Аргументы:
          x: входной тензор формы (batch_size, sequence_length, embedding_dim).

        Возвращает:
          Выходной тензор формы (batch_size, sequence_length, embedding_dim).
        """
        # Форма: (batch_size, sequence_length, embedding_dim).
        mlp_output = self.mlp(x)

        # Применяем dropout.
        mlp_output = self.dropout(mlp_output)

        return mlp_output

# Реализуйте здесь подсчёт числа параметров компонента MLP:
def parameter_count_mlp(hyperparams):

    parameter_count = ...

    return parameter_count


param_count = parameter_count_mlp(MODEL_HYPERPARAMETERS)
if param_count != ...:
    print(f"Number of parameters for MLP component: {param_count:,}")

# Запустите эту ячейку, чтобы проверить `parameter_count_mlp`.
feedback.test_parameter_count_mlp(parameter_count_mlp)

Задание 4. Блок трансформера

Следующая ячейка реализует блок трансформера: вход сначала проходит через механизм многоголового внимания, затем — через компонент MLP.

class TransformerBlock(layers.Layer):
  """Один блок трансформера.

    Блок трансформера — основной компонент архитектуры трансформера,
    которую обычно используют для задач с последовательностями. Он состоит из
    слоя MultiHeadAttention, за которым идёт MLP; на каждом шаге применяются
    нормализация слоя и dropout.

    Пример:
      transformer_block = TransformerBlock(embedding_dim=256, num_heads=8,
                                           mlp_dim=1024)
      output = transformer_block(inputs)

    Аргументы:
      embedding_dim: размерность входного эмбеддинга (и выхода слоя внимания).
      num_heads: число голов в многоголовом внимании.
      mlp_dim: число нейронов в MLP.
      dropout_rate: доля dropout, от 0 до 1.
      activation: функция активации в MLP.
    """

  def __init__(self,
               embedding_dim: int,
               num_heads: int,
               mlp_dim: int,
               dropout_rate: float = 0.0,
               activation: str = "relu",
               **kwargs: dict):
    super().__init__(**kwargs)

    self.self_attention = MultiHeadSelfAttention(embedding_dim,
                                                 num_heads,
                                                 dropout_rate)

    self.layer_norm_attention = LayerNorm(embedding_dim)
    self.feed_forward = MultiLayerPerceptron(embedding_dim,
                                             mlp_dim,
                                             dropout_rate,
                                             activation)
    self.layer_norm_mlp = LayerNorm(embedding_dim)

  def call(self, x: jax.Array) -> jax.Array:
    """Применяет один блок трансформера к входному тензору.

    Аргументы:
      x: входной тензор формы (batch_size, sequence_length, embedding_dim).

    Возвращает:
      Выходной тензор формы (batch_size, sequence_length, embedding_dim)
        после блока трансформера.
    """
    # Применяем маскированное self-attention.
    # Форма: (batch_size, sequence_length, embedding_dim).
    attention_output = self.self_attention(x)

    # Добавляем остаточную связь.
    attention_output = attention_output + x

    # Применяем нормализацию слоя (после сложения — схема Post-LN).
    attention_output = self.layer_norm_attention(attention_output)

    # Многослойный перцептрон применяется к выходу внимания.
    # Форма: (batch_size, sequence_length, embedding_dim).
    mlp_output = self.feed_forward(attention_output)

    # Добавляем остаточную связь.
    # Форма: (batch_size, sequence_length, embedding_dim).
    mlp_output = mlp_output + attention_output

    # Применяем нормализацию слоя.
    # Форма: (batch_size, sequence_length, embedding_dim).
    mlp_output = self.layer_norm_mlp(mlp_output)

    return mlp_output

# Реализуйте здесь подсчёт числа параметров одного блока трансформера:
def parameter_count_transformer_block(hyperparams):

    parameter_count = ...

    return parameter_count


param_count = parameter_count_transformer_block(MODEL_HYPERPARAMETERS)
if param_count != ...:
    print(f"Number of parameters for one transformer block: {param_count:,}")

# Запустите эту ячейку, чтобы проверить `parameter_count_transformer_block`.
feedback.test_parameter_count_transformer_block(parameter_count_transformer_block)

Задание 5. Выходной слой

Следующая ячейка реализует последний нужный компонент — выходной слой с функцией активации SoftMax. Этот слой выдаёт распределение вероятностей следующего токена.

class OutputLayer(keras.layers.Layer):
    """
    Слой, который вычисляет распределение логарифмов вероятностей по словарю.

    Проецирует входной тензор в размер словаря и применяет log-softmax.
    (Примечание переводчика: на деле метод call ниже возвращает сырые логиты,
    а SoftMax учитывается внутри функции потерь с from_logits=True.)

    Аргументы:
        vocabulary_size: размер словаря.
    """
    def __init__(self, vocabulary_size: int, **kwargs):
        super().__init__(**kwargs)
        self.vocabulary_size = vocabulary_size

        # Полносвязный слой проецирует вход из embedding_dim в vocabulary_size.
        self.output_layer = keras.layers.Dense(vocabulary_size,
                                        name="output_projection")

    def call(self, x: jax.Array) -> jax.Array:
        """
        Прямой проход OutputLayer.

        Аргументы:
          x: входной тензор формы (batch_size, sequence_length, embedding_dim).

        Возвращает:
          Логиты для каждого токена словаря, форма
            (batch_size, sequence_length, vocabulary_size).
        """
        # Проецируем измерение embedding_dim в размер словаря.
        logits = self.output_layer(x)

        return logits

# Реализуйте здесь подсчёт числа параметров выходного слоя:
def parameter_count_output_layer(hyperparams):

    parameter_count = ...

    return parameter_count


param_count = parameter_count_output_layer(MODEL_HYPERPARAMETERS)
if param_count != ...:
    print(f"Number of parameters for output layer: {param_count:,}")

# Запустите эту ячейку, чтобы проверить `parameter_count_output_layer`.
feedback.test_parameter_count_output_layer(parameter_count_output_layer)

Задание 6. Собираем всё вместе

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

входной слой эмбеддингов;
стопка из num_blocks блоков трансформера;
выходной слой.

class TransformerModel(layers.Layer):
    """Полный трансформер на Keras.

    Аргументы:
      vocabulary_size: размер словаря, то есть число уникальных токенов.
      max_length: максимальная длина входных последовательностей.
      embedding_dim: размерность пространства эмбеддингов.
      mlp_dim: число нейронов в MLP каждого блока трансформера.
      num_heads: число голов в многоголовом внимании.
      num_blocks: сколько блоков трансформера поставить друг за другом.
      dropout_rate: доля dropout для борьбы с переобучением.
      activation: функция активации в MLP каждого блока трансформера.
    """

    def __init__(self,
                 vocabulary_size: int,
                 max_length: int,
                 embedding_dim: int = 256,
                 mlp_dim: int = 256,
                 num_heads: int = 2,
                 num_blocks: int = 1,
                 dropout_rate: float = 0.0,
                 activation: str = "relu",
                 **kwargs):
        super().__init__(**kwargs)

        # Слой эмбеддингов: эмбеддинги токенов плюс позиционные.
        self.embedding_layer = TokenAndPositionEmbedding(vocabulary_size,
                                                         embedding_dim,
                                                         max_length)

        # Стопка блоков трансформера.
        self.transformer_blocks = keras.Sequential()
        for _ in range(num_blocks):
            self.transformer_blocks.add(TransformerBlock(
                embedding_dim,
                num_heads,
                mlp_dim,
                dropout_rate=dropout_rate,
                activation=activation))

        # Выходной слой.
        self.output_layer = OutputLayer(vocabulary_size)


    def call(self, x: jax.Array) -> jax.Array:

        # Эмбеддинги входных токенов.
        # Форма: (batch_size, sequence_length, embedding_dim).
        x = self.embedding_layer(x)

        # Форма: (batch_size, sequence_length, embedding_dim).
        x = self.transformer_blocks(x)

        # Выход (логиты для каждого токена словаря).
        # Форма: (batch_size, sequence_length, vocabulary_size).
        output = self.output_layer(x)

        return output

# Реализуйте здесь подсчёт общего числа параметров трансформера:
def parameter_count_transformer(hyperparams):

    parameter_count = ...

    return parameter_count


param_count = parameter_count_transformer(MODEL_HYPERPARAMETERS)
if param_count != ...:
    print(f"Number of parameters for full transformer model: {param_count:,}")

# Запустите эту ячейку, чтобы проверить `parameter_count_transformer`.
feedback.test_parameter_count_transformer(parameter_count_transformer)

Вы прошли по полной реализации трансформера: разобрали каждый компонент и то, как они соединяются.

Как вы увидели, у этой модели 135 миллионов обучаемых параметров — весьма немало. Для каждого из этих параметров оптимизатор должен вычислить градиенты на каждом примере и обновить веса на каждом шаге обучения. Поэтому обучение обычно занимает много времени. И заметьте: по сравнению с моделью вроде Gemma-1B эта модель всё ещё гораздо меньше.

Теперь вернитесь к числу параметров каждого компонента. Какие компоненты добавляют много параметров? Какие — мало?

Затем исследуйте, какие гиперпараметры сильно влияют на число параметров модели. Меняйте параметры ниже и смотрите, как меняются параметры каждого компонента.

# @title Считаем обучаемые параметры

max_length = 128  # @param {"type": "number"}
embedding_dim = 256  # @param {"type": "number"}
mlp_dim = 384  # @param {"type": "number"}
num_heads = 4  # @param {"type": "number"}
num_blocks = 2  # @param {"type": "number"}
vocabulary_size = 262144  # @param {"type": "number"}

your_hyperparameters = {
    "max_length": max_length,
    "embedding_dim": embedding_dim,
    "mlp_dim": mlp_dim,
    "num_heads": num_heads,
    "num_blocks": num_blocks,
    "vocabulary_size": vocabulary_size
}

# Считаем число параметров по компонентам.
parameter_counts = {
    "embedding": parameter_count_embedding(your_hyperparameters)
}

for i in range(your_hyperparameters["num_blocks"]):
    parameter_counts[f"transformer_block_{i}"] = parameter_count_transformer_block(your_hyperparameters)
    parameter_counts[f"  attention_{i}"] = parameter_count_attention(your_hyperparameters)
    parameter_counts[f"  layer_norm_attention_{i}"] = parameter_count_layer_norm(your_hyperparameters)
    parameter_counts[f"  mlp_{i}"] = parameter_count_mlp(your_hyperparameters)
    parameter_counts[f"  layer_norm_mlp_{i}"] = parameter_count_layer_norm(your_hyperparameters)

parameter_counts["output_layer"] = parameter_count_output_layer(your_hyperparameters)

# Печатаем таблицу числа параметров.
total = parameter_count_transformer(your_hyperparameters)
max_key_width = len(max(parameter_counts.keys(), key=len))
max_value_width = 80 - max_key_width

separator_length = max_key_width + 2 + max_value_width
print("Parameters of each component:")
separator_length = max_key_width + 2 + max_value_width
print("-" * separator_length)
print(f"{'Component':<{max_key_width}}  {'Parameters':>{max_value_width}}")
for key, value in parameter_counts.items():
    if not key.startswith(" "):
        print("-" * separator_length)
        print(f"{key:<{max_key_width}}  {value:>{max_value_width},}")
        if key.startswith("transformer_block"):
            print("~" * separator_length)
    else:
        print(f"{key:<{max_key_width}}  {value:>{max_value_width},}")

print("-" * separator_length)
print(f"{'Total':<{max_key_width}}  {total:>{max_value_width},}")

Ожидаемый вывод (наш расчёт по эталонным формулам) для гиперпараметров по умолчанию:

Компонент	Параметров	Как получается
embedding	67 108 864	$262\,144 \cdot 256$
attention (в каждом блоке)	263 168	$4 \cdot 257 \cdot 256$
layer_norm (каждая)	512	$2 \cdot 256$
mlp (в каждом блоке)	197 248	$257 \cdot 384 + 385 \cdot 256$
transformer_block (каждый)	460 928	$263\,168 + 197\,248 + 512$ (так считает эталон; точнее было бы $+\,1024$ , см. решение задания 4)
output_layer	67 371 008	$257 \cdot 262\,144$
Всего	135 401 728	$67\,108\,864 + 2 \cdot 460\,928 + 67\,371\,008$

Вывод из таблицы: 99,3% параметров этой модели — таблица эмбеддингов и выходной слой, то есть словарь на 262 144 токена. Два блока трансформера, где и происходит вся «умная» работа, — меньше одного процента. Уменьшите vocabulary_size до 3000 (как у вашего BPE) — и модель «похудеет» в сотни раз. Поэтому во многих моделях матрицы эмбеддингов и выходного слоя делают общими (weight tying): Gemma так и поступает.

Необязательно: обучаем модель

Последнее необязательное упражнение: если хотите посмотреть модель в деле, запустите следующую скрытую ячейку. Она загрузит набор Africa Galore, токенизирует и дополнит данные и обучит модель. Это займёт около минуты на экземпляре Colab с GPU или 10 минут на экземпляре с CPU.

Затем в ячейке после цикла обучения можно генерировать продолжения промпта.

Загружаем и токенизируем набор данных

Запустите следующую ячейку, чтобы загрузить набор Africa Galore и токенизировать его BPE-токенизатором из прошлых курсов.

# Загружаем набор данных и токенизатор.

africa_galore = pd.read_json(
    "https://storage.googleapis.com/dm-educational/assets/ai_foundations/africa_galore.json"
)
dataset = africa_galore["description"].values
print("Loaded dataset with", dataset.shape[0], "paragraphs.")
BPEWordTokenizer = tokenization.BPEWordTokenizer
tokenizer = BPEWordTokenizer.from_url("https://storage.googleapis.com/dm-educational/assets/ai_foundations/bpe_tokenizer_3000_v2.pkl")

encoded_tokens = []
for paragraph in tqdm.tqdm(dataset, unit="paragraphs"):
    encoded_tokens.append(tokenizer.encode(paragraph))

Готовим набор данных для обучения

Запустите следующую ячейку, чтобы дополнить и обрезать абзацы и подготовить входные и целевые последовательности.

max_length = 300
padded_sequences = keras.preprocessing.sequence.pad_sequences(
        encoded_tokens,
        maxlen=max_length,
        padding="post",
        truncating="post",
        value=tokenizer.pad_token_id,
    )
# Готовим вход и цель для трансформера.
# Для каждого примера берём все токены, кроме последнего.
input_sequences = padded_sequences[:, :-1]
# Для каждого примера берём все токены, кроме первого.
target_sequences = padded_sequences[:, 1:]

max_length = input_sequences.shape[1]

Задаём и обучаем модель

Следующая ячейка задаёт трансформер классом TransformerModel, который вы разобрали выше. Она также создаёт оптимизатор (Adam) и функцию потерь (многоклассовую перекрёстную энтропию) и прикрепляет оба компонента к модели. Наконец, она создаёт функцию, которая печатает генерацию каждые десять эпох, чтобы можно было следить за обучением.

# Фиксируем seed для воспроизводимости.
keras.utils.set_random_seed(3112)

# Создаём трансформер, который вы разобрали выше (маленький: E=64, M=128).
transformer = TransformerModel(
    tokenizer.vocabulary_size,
    max_length=max_length,
    dropout_rate=0.1,
    num_blocks=2,
    embedding_dim=64,
    mlp_dim=128,
)

# Собираем модель Keras.
input_layer = keras.Input((max_length,))
output_layer = transformer(input_layer)
model = keras.Model(input_layer, output_layer)

# Создаём оптимизатор.
optimizer = keras.optimizers.Adam(learning_rate=2.5e-3)
# Создаём функцию потерь.
loss_fn = keras.losses.SparseCategoricalCrossentropy(
    # Выходной слой выдаёт сырые логиты, а не вероятности после SoftMax:
    # так эффективнее и не возникает очень маленьких чисел.
    from_logits=True,
    ignore_class=tokenizer.pad_token_id,  # не учитываем <PAD> в потерях
    # Усредняем потери по батчу.
    reduction="sum_over_batch_size",
)

# Прикрепляем оптимизатор и функцию потерь к модели.
model.compile(optimizer=optimizer, loss=loss_fn)

# Обратный вызов, который печатает генерацию каждые 10 эпох.
prompt = "Jide"
prompt_ids = tokenizer.encode(prompt)
text_gen_callback = training.TextGenerator(
    max_tokens=11, start_tokens=prompt_ids, tokenizer=tokenizer, print_every=10
)

Запустите следующую ячейку, чтобы обучить модель методом fit 100 эпох.

model.fit(
    x=input_sequences,
    y=target_sequences,
    batch_size=32,
    epochs=100,
    callbacks=[text_gen_callback],
)

Промпт вашей модели

Наконец, следующей ячейкой можно давать промпты вашему трансформеру.

# Эта ячейка генерирует новые тексты.
# Чтобы поменять промпт, измените переменную prompt.
prompt = "Jide was thirsty so she went looking for a"

# Жадное сэмплирование.

generated_text, _ = generation.generate_text(
    prompt,
    n_tokens = 1,
    model = model,
    tokenizer = tokenizer,
    sampling_mode="greedy"
)
print(f"Generated text: {generated_text}")

Итоги

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

Лабораторная также показала, что разбиение модели на отдельные компоненты упрощает реализацию и поддержку кода сложных моделей. Вместо того чтобы реализовывать весь трансформер одной гигантской моделью Keras, можно собрать его из кубиков поменьше, каждый из которых проще. Такая модульность — важное свойство фреймворков вроде Keras.

Решения

Задание 1

def parameter_count_embedding(hyperparams: dict[str, int]) -> int:
    """Параметры матрицы эмбеддингов токенов: форма vocabulary_size x embedding_dim."""
    vocabulary_size = hyperparams["vocabulary_size"]
    embedding_dim = hyperparams["embedding_dim"]
    # Матрица эмбеддингов размера `vocabulary_size` x `embedding_dim`.
    parameter_count = vocabulary_size * embedding_dim
    return parameter_count

Задание 2

def parameter_count_attention(hyperparams: dict[str, int]) -> int:
    """Параметры подслоя многоголового внимания: четыре проекции Q, K, V и выхода,
    каждая — полносвязный слой embedding_dim x embedding_dim плюс смещения."""
    embedding_dim = hyperparams["embedding_dim"]
    # Для проекций ключей, запросов и значений выходная размерность —
    # d_head * num_heads, что как раз равно embedding_dim.
    q_parameter_count = (embedding_dim + 1) * embedding_dim  # проекция запросов
    k_parameter_count = (embedding_dim + 1) * embedding_dim  # проекция ключей
    v_parameter_count = (embedding_dim + 1) * embedding_dim  # проекция значений
    o_parameter_count = (embedding_dim + 1) * embedding_dim  # итоговая проекция

    parameter_count = (
        q_parameter_count
        + k_parameter_count
        + v_parameter_count
        + o_parameter_count
    )
    return parameter_count

Задание 3

def parameter_count_mlp(hyperparams: dict[str, int]) -> int:
    """Параметры MLP: два полносвязных слоя со смещениями,
    embedding_dim -> mlp_dim и mlp_dim -> embedding_dim."""
    embedding_dim = hyperparams["embedding_dim"]
    mlp_dim = hyperparams["mlp_dim"]

    # Первая проекция.
    ffn_parameter_count = (embedding_dim + 1) * mlp_dim
    # Вторая проекция.
    output_parameter_count = (mlp_dim + 1) * embedding_dim

    parameter_count = (
        ffn_parameter_count + output_parameter_count
    )
    return parameter_count

Задание 4

def parameter_count_transformer_block(hyperparams: dict[str, int]) -> int:
    """Параметры блока трансформера: внимание + MLP + две нормализации слоя."""
    embedding_dim = hyperparams["embedding_dim"]

    # Многоголовое внимание.
    mha_parameter_count = parameter_count_attention(hyperparams)

    # Компонент MLP.
    mlp_parameter_count = parameter_count_mlp(hyperparams)

    # Две нормализации слоя.
    layer_norm_parameter_count = 2 * embedding_dim

    parameter_count = (
        mha_parameter_count + mlp_parameter_count + layer_norm_parameter_count
    )
    return parameter_count

Внимание: в эталонном решении layer_norm_parameter_count = 2 * embedding_dim, то есть учтены как будто одна нормализация слоя ( $\gamma$ и $\beta$ по embedding_dim), хотя в блоке их две по $2 \cdot$ embedding_dim каждая. Точный счёт — 2 * parameter_count_layer_norm(hyperparams) = $4 \cdot 256 = 1024$ ; разница в 512 параметров на блок. Проверочная ячейка сверяется с эталоном, поэтому для зачёта оставьте формулу оригинала, а для себя запомните правильный ответ. (примечание переводчика)

Задание 5

def parameter_count_output_layer(hyperparams: dict[str, int]) -> int:
    """Параметры выходной проекции embedding_dim -> vocabulary_size со смещением
    для каждого токена словаря."""
    embedding_dim = hyperparams["embedding_dim"]
    vocabulary_size = hyperparams["vocabulary_size"]

    # Выходная проекция.
    output_parameter_count = (embedding_dim + 1) * vocabulary_size

    # Параметры есть только у проекции, у функции активации их нет.
    parameter_count = output_parameter_count

    return parameter_count

Задание 6

def parameter_count_transformer(hyperparams: dict[str, int]) -> int:
    """Параметры всего трансформера из `num_blocks` блоков."""
    num_blocks = hyperparams["num_blocks"]

    # Слой эмбеддингов.
    embedding_parameter_count = parameter_count_embedding(hyperparams)

    # `num_blocks` блоков трансформера.
    transformer_blocks_parameter_count = (
        num_blocks * parameter_count_transformer_block(hyperparams)
    )

    # Выходной слой.
    output_parameter_count = parameter_count_output_layer(hyperparams)

    parameter_count = (
        embedding_parameter_count
        + transformer_blocks_parameter_count
        + output_parameter_count
    )

    return parameter_count

Перевод выполнен по лицензии оригинала (CC-BY-4.0 (тексты), Apache-2.0 (код)); автор оригинала — Google DeepMind и UCL. Нашли неточность — сверяйтесь с оригиналом по ссылке выше. · обновлено 2026-10-04