AI Research Foundations, курс 4, часть 1: внимание своими руками

Перевод лабораторных 4.1–4.3 курса Google DeepMind «Discover the Transformer Architecture»: карты внимания Gemma-1B по слоям, формула внимания на настоящих матрицах Q, K, V и маскированное многоголовое внимание.

Полный переводс языка: английский

Оригинал: AI Research Foundations — Course 4: Discover the Transformer Architecture, Labs 4.1–4.3 · 3 ноутбука Colab, ≈50 минут · ветка main, коммит a3b4aa7 (2026-09-02)
Автор: Google DeepMind и UCL
Лицензия оригинала: CC-BY-4.0 (тексты), Apache-2.0 (код)

Текст оригинала переведён целиком, его лицензия это разрешает.

К урокам курса

Курс 4 программы AI Research Foundations называется «Discover the Transformer Architecture» («Откройте архитектуру трансформера»). В курсах 1–2 трансформер был готовым кубиком из пакета курса; теперь вы разбираете его главный механизм — внимание (attention). Сначала вы посмотрите, на какие токены «смотрит» настоящая модель Gemma-1B, потом сами посчитаете внимание по её матрицам запросов, ключей и значений и добавите маску и несколько голов. Продолжение — часть 2: позиционные эмбеддинги и параметры трансформера.

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

Лабораторная	Время	Когда делать	Что нужно
4.1. Визуализация весов внимания	10 мин	Вместе с уроком «Внимание интуитивно»	Gemma-1B, желательно GPU
4.2. Реализуйте формулу внимания	15 мин	Вместе с уроком «Scaled dot-product self-attention»	JAX, Gemma-1B
4.3. Маскированное многоголовое внимание	25 мин	Вместе с уроком «Многоголовое внимание»	JAX, Gemma-1B

В уроках ChinaAI вы пишете внимание на PyTorch для своей маленькой модели; здесь та же формула проверяется на весах настоящей Gemma-1B — ваш код должен дать картинку, совпадающую с эталонной. К лабораторной 4.1 хорошо вернуться на ступени 9, в уроке «Заглядываем внутрь трансформера»: карты внимания — первый инструмент интерпретируемости.

Лабораторная 4.1. Визуализация весов внимания

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

Исследуйте, как веса внимания трансформера показывают, какие токены влияют на предсказание. 10 минут.

Обзор

В этой лабораторной вы разберётесь, как трансформеры делают предсказания. Трансформеры во многом опираются на так называемый механизм внимания (attention mechanism). Он позволяет модели строить контекстные эмбеддинги (contextual embeddings), объединяющие информацию от всех токенов промпта. При обучении трансформера этот механизм выучивает, сколько внимания уделять каждому токену промпта, то есть какой вес придавать информации каждого токена, когда строится контекстный эмбеддинг для предсказания следующего токена.

В отличие от упражнения с ручкой и бумагой в начале курса, трансформер решает, какие токены важны, не один раз. У него много слоёв (layers), и механизм внимания работает в каждом из них. Например, у модели Gemma-1B 26 слоёв.

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

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

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

Задачи

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

изобразите веса внимания в разных слоях модели Gemma-1B для разных промптов;
исследуете веса внимания для промптов из нескольких предложений.

Если есть возможность, советуем запускать лабораторную на экземпляре Colab с GPU.

Импорты

В этой лабораторной вы используете пакет курса ai_foundations: в нём есть методы, чтобы загрузить Gemma-1B и изобразить её веса внимания.

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

from ai_foundations import generation  # Чтобы загрузить Gemma и давать ей промпты.
from ai_foundations import visualizations  # Чтобы рисовать веса внимания.

Загружаем модель

Запустите следующую ячейку, чтобы загрузить модель Gemma-1B. Этой моделью вы будете генерировать следующий токен и вычислять её веса внимания. Ячейка также создаёт кэш, который ускоряет построение картинок, если вы меняете только номер слоя.

# Создаём кэши для промптов.
previous_prompt = None
previous_prompt2 = None

# Загружаем особую версию Gemma-1B, которая даёт доступ к весам внимания.
print("Loading Gemma-1B...")
model = generation.load_gemma("Gemma-1B-AttentionWeight")
print("Loaded Gemma-1B.")

Веса внимания для промптов из одного предложения

Введите промпт ниже и запустите ячейку, чтобы изобразить веса внимания для конкретного слоя.
Исследуйте, как веса меняются от слоя к слою. Смотрит ли модель на одни и те же токены в каждом слое? Совпадают ли веса внимания с тем, что, по-вашему, важно для предсказания?
Исследуйте, как меняются веса, если добавить слова в промпт или убрать их.

Код ниже строит рисунок, который показывает, насколько представление каждого токена промпта важно для предсказания сгенерированного токена. Чем толще линия, тем сильнее представление токена влияет на предсказание следующего токена в этом слое.

Учтите, что трансформерные языковые модели часто придают большой вес первому токену — специальному токену начала последовательности (BOS) (см., например, статью Analyzing the Structure of Attention in a Transformer Language Model [1]). Поэтому на картинках ниже линии к токену BOS (<BOS>) серые — так легче разглядеть остальные связи.

# @title Рисуем веса внимания
layer = 19  # @param {type:"slider", min: 0, max: 25}

prompt = "Jide was hungry so she went looking for"  # @param {type: "string"}

show_all_weights = False  # показывать веса только для последнего (сгенерированного) токена


if prompt != previous_prompt:
    # Пересчитываем, только если промпт изменился (кэш).
    (output_text, next_token_logits, tokenizer, attention_weights, _, _) = (
        generation.prompt_attention_transformer_model(
            prompt, model, sampling_mode="greedy"
        )
    )
    tokens = [tokenizer.tokens[t] for t in tokenizer.encode(output_text)]
    previous_prompt = prompt

print(f"Generated text: {output_text}")

visualizations.visualize_attention(
    tokens,
    attention_weights[f"layer_{layer}"],
    layer,
    min_line_thickness=0,
    max_line_thickness=5,
    show_all_weights=show_all_weights
)

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

Для промпта «Jide was hungry so she went looking for» модель предсказала самым вероятным следующим токеном «food». В некоторых слоях (например, в 19-м) это предсказание сильно зависит от токена «hungry» («голодная»). В других (например, во 2-м) — от совсем локального контекста, например «looking for» («искать»).

Меняя промпт, вы, вероятно, видели похожие закономерности. Например, заменив промпт на «Jide was thirsty so she went looking for», вы, скорее всего, увидели, что в 19-м слое модель для предсказания следующего токена «water» сильно опирается на токен «thirsty» («хочет пить»).

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

Внимание через границы предложений

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

Введите ниже промпт из нескольких предложений. Например, можно скопировать короткий текст из предыдущего упражнения:

«The artist selected a shade of blue from her palette. With a steady hand, she began to paint the waves onto the canvas. She wanted to capture the power and beauty of the» («Художница выбрала на палитре оттенок синего. Твёрдой рукой она начала писать волны на холсте. Она хотела передать силу и красоту…»)
Исследуйте, какие токены важнее всего для предсказания. Все ли они из одного предложения?
Сравните веса для этого промпта с весами для промпта из одного предложения:
- в каких слоях поведение похоже, несмотря на разные тексты?
- найдите слой, сосредоточенный на смысле. Какой вес он даёт токенам, важнее всего для понимания контекста?
- теперь найдите слой, сосредоточенный на грамматике. Какой вес он даёт токенам, важнее всего для грамматической правильности предложения?

Примечание: если картинку трудно разглядеть, щёлкните по ней, чтобы увеличить.

# @title Рисуем веса внимания
layer = 19  # @param {type:"slider", min: 0, max: 25}

prompt = "The artist selected a shade of blue from her palette. With a steady hand, she began to paint the waves onto the canvas. She wanted to capture the power and beauty of the"  # @param {type: 'string'}
show_all_weights = False


if prompt != previous_prompt2:
    (output_text2, next_token_logits2, tokenizer, attention_weights2, _, _) = (
        generation.prompt_attention_transformer_model(
            prompt, model, sampling_mode="greedy"
        )
    )

    tokens2 = [tokenizer.tokens[t] for t in tokenizer.encode(output_text2)]
    previous_prompt2 = prompt

print(f"Generated text: {output_text2}")

visualizations.visualize_attention(
    tokens2,
    attention_weights2[f"layer_{layer}"],
    layer,
    min_line_thickness=0,
    max_line_thickness=5,
    show_all_weights=show_all_weights,
)

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

Итоги

В этом коротком задании вы изобразили веса внимания модели Gemma-1B для разных промптов. Вы увидели, что одни токены важнее для предсказания следующего токена, чем другие, и что этот набор токенов может меняться от слоя к слою. Вы также поэкспериментировали с длинными промптами и увидели, что предсказание следующего токена может зависеть и от токенов предыдущих предложений.

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

Литература

[1] Jesse Vig and Yonatan Belinkov. 2019. Analyzing the Structure of Attention in a Transformer Language Model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP. https://aclanthology.org/W19-4808/.

Лабораторная 4.2. Реализуйте формулу внимания

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

Вычислите вектор весов внимания для всех слоёв трансформера. 15 минут.

Обзор

В этой лабораторной вы сами реализуете формулу внимания. Вы возьмёте выученные матрицы проекций $W_Q$ , $W_K$ , $W_V$ модели Gemma-1B и по этим трём матрицам вычислите веса внимания и новый контекстный эмбеддинг.

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

К концу лабораторной вы поймёте, как веса внимания вычисляются из матриц запросов (query), ключей (key) и значений (value).

Задачи

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

загрузите обученную модель Gemma-1B и вычислите ею матрицы запросов, ключей и значений;
реализуете формулу внимания;
изобразите свои веса внимания и сравните их с весами эталонной реализации.

Если есть возможность, советуем запускать лабораторную на экземпляре Colab с GPU.

Импорты

В этой лабораторной вы в основном работаете с пакетом JAX, чтобы менять векторы и матрицы. Функции пакета курса ai_foundations загружают Gemma, достают её параметры и строят картинки.

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

import jax  # Для работы с векторами и матрицами.
import jax.numpy as jnp  # Для работы с векторами и матрицами.
from ai_foundations import generation  # Чтобы давать промпты модели Gemma.
from ai_foundations import visualizations  # Чтобы рисовать веса внимания.
from ai_foundations import attention  # Чтобы работать с матрицами Q, K, V.
# Для обратной связи по вашей реализации.
from ai_foundations.feedback.course_4 import attention as attention_feedback

Загружаем модель и рисуем веса внимания

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

# Создаём кэши для промптов.
previous_prompt = None
previous_prompt2 = None

# Загружаем особую версию Gemma-1B, которая даёт доступ к весам внимания
# и матрицам Q, K, V.
print("Loading Gemma-1B...")
model = generation.load_gemma("Gemma-1B-AttentionWeight")
print("Loaded Gemma-1B.")

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

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

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

# @title Рисуем веса внимания (эталонная реализация)
layer = 19  # @param {type:"slider", min: 0, max: 25}

prompt = "Jide was hungry so she went looking for"  # @param {type: "string"}
# @markdown Отметьте, чтобы показать веса внимания для всех токенов, а не только для сгенерированного:
show_all_weights = True  # @param {type:"boolean"}


if prompt != previous_prompt:
  (
      output_text,
      _,
      tokenizer,
      attention_weights,
      _,
      qkv_dict,  # проекции Q, K, V всех слоёв — понадобятся в задании
  ) = generation.prompt_attention_transformer_model(
      prompt, model, sampling_mode="greedy"
  )
  tokens = [tokenizer.tokens[t] for t in tokenizer.encode(output_text)]
  previous_prompt = prompt

print(f"Generated text: {output_text}")

visualizations.visualize_attention(
    tokens,
    attention_weights[f"layer_{layer}"],
    layer,
    min_line_thickness=0,
    max_line_thickness=5,
    show_all_weights=show_all_weights,
)

Задание 1. Формула внимания

Реализуйте в следующей ячейке механизм внимания для конкретного слоя. Для простоты ячейка уже вызывает get_qkv_matrices, которая возвращает проекции запросов, ключей и значений для каждого токена предложения.

Изучите матрицы проекций запросов, ключей и значений. Какой у них размер?
Вычислите $d_k$ — размерность векторов-ключей; она нужна для нормирующего множителя. Для этого можно использовать .shape.
Вычислите логиты: перемножьте матрицы проекций запросов и ключей и поделите на квадратный корень из $d_k$ :
$\text{logits} = \frac{QK^T}{\sqrt{d_k}}$
Вместо jnp.matmul можно использовать оператор @: например, A @ B вычисляет произведение матриц A и B. Квадратный корень — функция jnp.sqrt. И помните, что матрицу A можно транспонировать, вызвав A.T.
Вычислите веса внимания alpha, применив к логитам функцию SoftMax. Можно написать свою SoftMax или использовать jax.nn.softmax:
$\alpha = \text{SoftMax}(\text{logits})$
Вычислите выход механизма внимания, умножив веса внимания $\alpha$ на проекции значений:
$Y = \alpha V$

Дописав недостающие части следующей ячейки, запустите проверочную ячейку и ячейку «Рисуем веса внимания (ваша реализация)», чтобы увидеть веса внимания.

Как читать формулу словами: строка $i$ матрицы $QK^T$ — скалярные произведения запроса токена $i$ с ключами всех токенов, то есть «насколько каждый токен подходит под то, что ищет токен $i$ ». Деление на $\sqrt{d_k}$ не даёт числам разрастаться с ростом размерности, SoftMax превращает строку в доли, которые в сумме дают 1, а умножение на $V$ смешивает векторы-значения токенов в этих долях. Пример с числами: если для двух токенов логиты строки равны $(2, 0)$ , SoftMax даёт веса $(0{,}88;\ 0{,}12)$ , и выход — это $0{,}88 \cdot v_1 + 0{,}12 \cdot v_2$ .

def compute_attention(
    qkv_dict: dict[str, dict[str, jax.Array]], layer: int
) -> tuple[jax.Array, jax.Array, jax.Array]:
  """Вычисляет веса внимания для одного слоя.

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

  Возвращает:
    Y: выход слоя внимания. Форма: (n_tokens, embedding_dim).
    alpha: веса внимания. Форма: (n_tokens, n_tokens).
    logits: исходные логиты, из которых получаются веса внимания.
      Форма: (n_tokens, n_tokens).
  """
  # Достаём из `qkv_dict` матрицы проекций запросов, ключей и значений.
  query_proj, key_proj, value_proj = attention.get_qkv_matrices(qkv_dict, layer)

  logits = ...  # Вычислите логиты здесь.

  alpha = ...  # Вычислите веса внимания здесь.

  Y = ...  # Вычислите выход здесь.

  return Y, alpha, logits

# Запустите эту ячейку, чтобы проверить реализацию.
attention_feedback.test_compute_attention(compute_attention, qkv_dict)

Рисуем ваши веса внимания

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

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

# @title Рисуем веса внимания (ваша реализация)
layer = 19  # @param {type:"slider", min: 0, max: 25}

prompt = "Jide was hungry so she went looking for"  # @param {type: 'string'}
# @markdown Отметьте, чтобы показать веса внимания для всех токенов, а не только для сгенерированного:
show_all_weights = True  # @param {type:"boolean", label: "sfdsf"}

model_name = "Gemma-1B"

if prompt != previous_prompt2:
  (
      output_text2,
      _,
      tokenizer,
      attention_weights2,
      attention_mask2,
      qkv_dict2,
  ) = generation.prompt_attention_transformer_model(
      prompt, model, sampling_mode="greedy"
  )
  tokens2 = [tokenizer.tokens[t] for t in tokenizer.encode(output_text2)]
  previous_prompt2 = prompt

print(f"Generated text: {output_text2}")


attn_out, attn_weights, attn_logits = compute_attention(qkv_dict2, layer)
# Дополнительно обрабатываем веса внимания — применяем маску (о ней в следующей
# лабораторной), чтобы веса совпали с теми, что вычисляет Gemma.
attn_weights_post_processed = attention.post_process_attention(
    attn_logits, attention_mask2
)

visualizations.visualize_attention(
    tokens2,
    attn_weights_post_processed,
    layer,
    min_line_thickness=0,
    max_line_thickness=5,
    show_all_weights=show_all_weights,
)

Итоги

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

Решения

Задание 1

def compute_attention(
    qkv_dict: dict[str, dict[str, jax.Array]], layer: int
) -> tuple[jax.Array, jax.Array, jax.Array]:
  # Достаём из `qkv_dict` матрицы проекций запросов, ключей и значений.
  query_proj, key_proj, value_proj = attention.get_qkv_matrices(qkv_dict, layer)

  # Достаём d_k: форма key_proj — (число токенов, d_k).
  _, dim_key = key_proj.shape

  # Вычисляем logits = QK^T / sqrt(d_k).
  logits = (query_proj @ key_proj.T) / jnp.sqrt(dim_key)

  # Вычисляем веса внимания 𝛼 (SoftMax по каждой строке).
  alpha = jax.nn.softmax(logits)

  # Вычисляем выход Y = 𝛼V.
  Y = alpha @ value_proj

  return Y, alpha, logits

Лабораторная 4.3. Маскированное многоголовое внимание

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

Вычислите веса внимания для нескольких голов и примените маску, чтобы модель не «заглядывала в будущее». 25 минут.

Обзор

В этой лабораторной вы реализуете маскированное многоголовое внимание (masked multi-head attention). Как и в прошлой лабораторной, вы возьмёте выученные матрицы проекций $W_Q$ , $W_K$ , $W_V$ модели Gemma-1B и по ним вычислите веса внимания и выход механизма внимания.

Для простоты в двух прошлых лабораториях рассматривалась только первая голова внимания (attention head) модели Gemma-1B. В этой лабораторной можно изобразить и остальные головы. У Gemma-1B в каждом слое четыре головы внимания.

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

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

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

Задачи

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

реализуете функцию, которая создаёт маску внимания;
реализуете маскированное многоголовое внимание;
изобразите свои веса внимания и сравните их с весами эталонной реализации.

flowchart LR
QKV["Q, K, V по головам<br/>(список матриц)"] --> S["stack_matrices<br/>(головы, токены, d)"]
S --> L["compute_raw_logits<br/>QK^T / sqrt(d_k)"]
L --> M["apply_attention_mask<br/>будущее → K_MASK"]
M --> A["compute_attention_weights<br/>SoftMax"]
A --> Y["compute_attention_output<br/>Y = alpha V"]

Шесть функций лабораторной 4.3 и порядок, в котором их вызывает compute_attention

Импорты

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

import jax  # Для работы с векторами и матрицами.
import jax.numpy as jnp  # Для работы с векторами и матрицами.
from ai_foundations import generation  # Чтобы давать промпты модели Gemma.
from ai_foundations import visualizations  # Чтобы рисовать веса внимания.
from ai_foundations import attention  # Чтобы работать с матрицами Q, K, V.
# Для обратной связи по вашей реализации.
from ai_foundations.feedback.course_4 import attention as attention_feedback

Загружаем модель и рисуем веса внимания

Вы возьмёте параметры модели Gemma-1B, чтобы вычислять веса внимания. Запустите следующую ячейку, чтобы загрузить параметры Gemma-1B и создать кэши.

# Создаём кэши для промптов.
previous_prompt = None
previous_prompt2 = None

# Загружаем особую версию Gemma-1B, которая даёт доступ к весам внимания
# и матрицам Q, K, V.
print("Loading Gemma-1B...")
model = generation.load_gemma("Gemma-1B-AttentionWeight")
print("Loaded Gemma-1B.")

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

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

# @title Рисуем веса внимания (эталонная реализация)
from IPython.display import clear_output

layer = 19  # @param {type:"slider", min: 0, max: 25}

head = 3  # @param {type:"slider", min: 0, max: 3}

prompt = "Jide was hungry so she went looking for"  # @param {type: "string"}
# @markdown Отметьте, чтобы показать веса внимания для всех токенов, а не только для сгенерированного:
show_all_weights = True  # @param {type:"boolean"}

if prompt != previous_prompt:
    (
        output_text,
        next_token_logits,
        tokenizer,
        attention_weights,
        _,
        qkv_dict,
    ) = generation.prompt_attention_transformer_model(
        prompt, model, sampling_mode="greedy"
    )
    tokens = [tokenizer.tokens[t] for t in tokenizer.encode(output_text)]
    previous_prompt = prompt

print(f"Generated text: {output_text}")

visualizations.visualize_attention(
    tokens,
    attention_weights[f"layer_{layer}"],
    layer,
    head=head,
    min_line_thickness=0,
    max_line_thickness=5,
    show_all_weights=show_all_weights,
)

Задание 1. Маска внимания

def compute_attention_mask(num_tokens: int) -> jax.Array:
    """Вычисляет маску внимания для входа длиной `num_tokens`.

    Аргументы:
      num_tokens: число токенов во входе.

    Возвращает:
      attention_mask: треугольная матрица: какие логиты внимания оставить (1),
        а какие замаскировать (0). Форма: (num_tokens, num_tokens).
    """

    attention_mask = ...  # Вычислите маску внимания здесь.

    return attention_mask

print(f"Attention mask for input length 5:\n {compute_attention_mask(5)}")
print(f"\nAttention mask for input length 6:\n {compute_attention_mask(6)}")

# Запустите эту ячейку, чтобы проверить реализацию `compute_attention_mask`.
attention_feedback.test_compute_attention_mask(compute_attention_mask)

Ожидаемый вывод для длины 5 (jnp.tri возвращает числа с плавающей точкой):

[[1. 0. 0. 0. 0.]
 [1. 1. 0. 0. 0.]
 [1. 1. 1. 0. 0.]
 [1. 1. 1. 1. 0.]
 [1. 1. 1. 1. 1.]]

Строка $i$ — токен, который «смотрит», столбец $j$ — токен, на который смотрят. Единицы только на диагонали и ниже: токен видит себя и всё, что было до него, но не видит будущего.

Реализуем многоголовое внимание

В следующих пяти заданиях вы шаг за шагом реализуете механизм многоголового внимания.

Задание 2. Складываем матрицы в стопку

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

Проекции запросов, ключей и значений хранятся в списках, где каждый элемент — матрица проекции одной головы (например, первые элементы всех трёх списков — матрицы $Q$ , $K$ и $V$ первой головы внимания). stack_matrices должна сложить эти матрицы в стопку так, чтобы для $Q$ , $K$ и $V$ получились трёхмерные тензоры (то есть матрицы, но с тремя измерениями). Первое измерение — голова, второе — номер токена, третье — компоненты векторов. Для этого удобна функция jnp.stack.

Допишите реализацию в следующей ячейке и проверьте её, запустив ячейку после неё.

def stack_matrices(
    qkv_dict: dict[str, dict[str, jax.Array]], layer: int
) -> tuple[jax.Array, jax.Array, jax.Array]:
    """Достаёт списки матриц запросов, ключей и значений для каждой головы
        и складывает их в многомерные тензоры.

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

    Возвращает:
      query_proj: тензор проекций запросов.
        Форма: (num_heads, num_tokens, dim_k).
      key_proj: тензор проекций ключей. Форма: (num_heads, num_tokens, dim_k).
      value_proj: тензор проекций значений.
        Форма: (num_heads, num_tokens, dim_v).
    """

    # Если передать `head=None`, `get_qkv_matrices` вернёт списки из
    # `num_heads` матриц проекций, а не одну матрицу.
    (
        query_proj_list,
        key_proj_list,
        value_proj_list
    ) = attention.get_qkv_matrices(
        qkv_dict, layer, head=None
    )

    # Сложите матрицы проекций запросов, ключей и значений в стопку здесь:
    query_proj = ...
    key_proj = ...
    value_proj = ...

    return query_proj, key_proj, value_proj

# Запустите эту ячейку, чтобы проверить реализацию `stack_matrices`.
attention_feedback.test_stack_matrices(stack_matrices, qkv_dict)

Если заглянуть в код attention.get_qkv_matrices, видно, что у каждого слоя Gemma 3 1B четыре разные матрицы запросов, но одна общая матрица ключей и одна — значений: функция просто копирует (jnp.broadcast_to) её для каждой из четырёх голов. Это внимание с одной общей головой ключей и значений (multi-query attention, MQA) — частный случай GQA. Оно экономит память KV-кэша при генерации; подробно — в уроке «MQA и GQA». Для вашего кода ничего не меняется: после копирования это обычные тензоры формы (4, число токенов, 256).

Задание 3. Вычисляем исходные логиты

Теперь допишите функцию compute_raw_logits, которая вычисляет логиты по матрицам запросов и ключей.

Функция должна выполнить три операции:

Вычислить $d_k$ — размерность векторов-ключей для нормирующего множителя. Помните, что для этого можно использовать .shape.
Транспонировать матрицу проекций ключей, $K^T$ . Это уже сделано.
Вычислить логиты: перемножить проекции запросов и ключей и поделить на квадратный корень из $d_k$ :

\text{logits} = \frac{QK^T}{\sqrt{d_k}}

Допишите реализацию в следующей ячейке и проверьте её, запустив ячейку после неё.

def compute_raw_logits(query_proj: jax.Array, key_proj: jax.Array) -> jax.Array:
    """Вычисляет исходные логиты по тензорам проекций запросов и ключей —
    до применения маски внимания.

    Аргументы:
      query_proj: тензор проекций запросов.
        Форма: (num_heads, num_tokens, dim_k).
      key_proj: тензор проекций ключей.
        Форма: (num_heads, num_tokens, dim_k).

    Возвращает:
      logits: исходные логиты. Форма: (num_heads, num_tokens, num_tokens).
    """

    # Меняем местами два последних измерения тензора ключей, чтобы их можно было
    # перемножить. Форма меняется с (num_heads, num_tokens, dim_k) на
    # (num_heads, dim_k, num_tokens).
    key_transposed = jnp.transpose(key_proj, (0, 2, 1))

    logits = ...  # Вычислите исходные логиты здесь.

    return logits

# Запустите эту ячейку, чтобы проверить реализацию `compute_raw_logits`.
attention_feedback.test_compute_raw_logits(compute_raw_logits, qkv_dict)

Оператор @ для трёхмерных тензоров перемножает матрицы «по головам»: тензор (4, n, 256) на (4, 256, n) даёт (4, n, n) — четыре отдельные матрицы логитов, по одной на голову.

Задание 4. Применяем маску внимания

Допишите функцию apply_attention_mask, которая вычисляет маску внимания и применяет её к исходным логитам из прошлого шага.

Функция должна выполнить две операции:

Вычислить маску внимания вашей функцией compute_attention_mask.
Вычислить замаскированные логиты:
$\text{logits\_masked}_{i,j} = \begin{cases} \text{logits}_{i,j}, & \text{если } \text{mask}_{i,j} = 1, \\ K & \text{иначе.} \end{cases}$
Чтобы вычислить функцию, зависящую от такого условия, используйте jnp.where. $K$ задано в следующей ячейке как K_MASK.

Допишите реализацию в следующей ячейке и проверьте её, запустив ячейку после неё.

K_MASK = -2.3819763e+38  # очень большое по модулю отрицательное число: exp(K_MASK) = 0

def apply_attention_mask(logits_raw: jax.Array) -> jax.Array:
    """Вычисляет маску внимания (через `compute_attention_mask`) и
    применяет её к логитам.

    Аргументы:
      logits_raw: исходные логиты. Форма: (num_heads, num_tokens, num_tokens).

    Возвращает:
      Замаскированные логиты. Форма: (num_heads, num_tokens, num_tokens).
    """

    _, num_tokens, _ = logits_raw.shape

    attention_mask = ...  # Вычислите маску внимания здесь.
    logits_masked = ...  # Вычислите замаскированные логиты здесь.

    return logits_masked

# Запустите эту ячейку, чтобы проверить реализацию `apply_attention_mask`.
attention_feedback.test_apply_attention_mask(apply_attention_mask, qkv_dict)

Почему огромное отрицательное число, а не просто ноль в весах? Маска применяется до SoftMax: $e^{-2{,}38 \cdot 10^{38}}$ равно нулю, поэтому будущие токены получают вес ровно 0, а веса остальных по-прежнему в сумме дают 1. Если бы мы обнулили веса после SoftMax, сумма строки стала бы меньше 1.

Задание 5. Вычисляем веса внимания

Допишите функцию compute_attention_weights, которая вычисляет веса внимания alpha.

alpha получается, если применить к логитам функцию SoftMax. Можно написать свою SoftMax или использовать jax.nn.softmax:

\alpha = \text{SoftMax}(\text{logits})

Учтите, что по умолчанию jax.nn.softmax нормирует значения многомерной матрицы или тензора так, чтобы в сумме давали 1 значения вдоль последнего измерения. Здесь это правильное поведение, но в других случаях может понадобиться указать другое измерение аргументом axis.

Допишите реализацию в следующей ячейке и проверьте её, запустив ячейку после неё.

def compute_attention_weights(logits_masked: jax.Array) -> jax.Array:
    """Вычисляет веса внимания `alpha` из замаскированных логитов,
    применяя функцию SoftMax.

    Аргументы:
      logits_masked: замаскированные логиты из прошлых шагов.
        Форма: (num_heads, num_tokens, num_tokens).

    Возвращает:
      alpha: веса внимания. Форма: (num_heads, num_tokens, num_tokens).
    """

    alpha = ...  # Примените SoftMax здесь.

    return alpha

# Запустите эту ячейку, чтобы проверить реализацию `compute_attention_weights`.
attention_feedback.test_compute_attention_weights(compute_attention_weights, qkv_dict)

Задание 6. Вычисляем выход внимания

def compute_attention_output(
    alpha: jax.Array, value_proj: jax.Array
) -> jax.Array:
    """Вычисляет выход механизма внимания: взвешенную сумму эмбеддингов
    из `value_proj` с весами внимания `alpha`.

    Аргументы:
      alpha: веса внимания. Форма: (num_heads, num_tokens, num_tokens).
      value_proj: тензор проекций значений.
        Форма: (num_heads, num_tokens, dim_v).

    Возвращает:
      Y: выход механизма внимания.
        Форма: (num_heads, num_tokens, dim_v).
    """

    Y = ...  # Вычислите выход механизма внимания.

    return Y

# Запустите эту ячейку, чтобы проверить реализацию `compute_attention_output`.
attention_feedback.test_compute_attention_output(compute_attention_output, qkv_dict)

Собираем всё вместе

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

Обратите внимание: в этом задании вы пропускаете шаг проекции, на котором выходы всех голов внимания объединяются в итоговый выход многоголового внимания. Вместо этого compute_attention возвращает выход каждой головы по отдельности.

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

def compute_attention(
    qkv_dict: dict[str, dict[str, jax.Array]], layer: int
) -> tuple[jax.Array, jax.Array]:
    """Реализует многоголовое внимание с помощью функций выше:
      - складывает матрицы проекций в стопку;
      - вычисляет исходные и замаскированные логиты;
      - вычисляет веса внимания alpha;
      - вычисляет выход внимания Y.

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

    Возвращает:
      Y: выход многоголового внимания.
        Форма: (num_tokens, num_heads, dim_v).
      alpha: веса внимания каждой головы.
        Форма: (num_tokens, num_heads, num_tokens).
    """

    # Складываем матрицы в стопку.
    # Формы:
    #   query_proj: (num_heads, num_tokens, dim_k).
    #   key_proj: (num_heads, num_tokens, dim_k).
    #   value_proj: (num_heads, num_tokens, dim_v).
    query_proj, key_proj, value_proj = stack_matrices(qkv_dict, layer)

    # Исходные логиты. Форма: (num_heads, num_tokens, num_tokens).
    logits_raw = compute_raw_logits(query_proj, key_proj)

    # Замаскированные логиты. Форма: (num_heads, num_tokens, num_tokens).
    logits_masked = apply_attention_mask(logits_raw)

    # Веса внимания. Форма: (num_heads, num_tokens, num_tokens).
    alpha = compute_attention_weights(logits_masked)

    # Выход механизма внимания.
    # Форма: (num_heads, num_tokens, dim_v).
    Y = compute_attention_output(alpha, value_proj)

    # Переставляем измерения так, чтобы первым был номер токена, а вторым —
    # голова: так Gemma хранит матрицы весов внимания.
    # Форма: (num_tokens, num_heads, num_tokens).
    alpha = jnp.transpose(alpha, (1, 0, 2))
    Y = jnp.transpose(Y, (1, 0, 2))  # Форма: (num_tokens, num_heads, dim_v).

    return Y, alpha

Рисуем ваши веса внимания

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

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

# @title Рисуем веса внимания (ваша реализация)
layer = 19  # @param {type:"slider", min: 0, max: 25}
head = 3  # @param {type:"slider", min: 0, max: 3}

prompt = "Jide was hungry so she went looking for"  # @param {type: 'string'}
# @markdown Отметьте, чтобы показать веса внимания для всех токенов, а не только для сгенерированного:
show_all_weights = True  # @param {type:"boolean", label: "sfdsf"}

model_name = "Gemma-1B"

if prompt != previous_prompt2:
    (
        output_text2,
        next_token_logits2,
        tokenizer,
        attention_weights2,
        attention_mask2,
        qkv_dict2,
    ) = generation.prompt_attention_transformer_model(
        prompt, model, sampling_mode="greedy"
    )
    tokens2 = [tokenizer.tokens[t] for t in tokenizer.encode(output_text2)]
    previous_prompt2 = prompt

print(f"Generated text: {output_text2}")

attn_out, attn_weights = compute_attention(qkv_dict2, layer)

visualizations.visualize_attention(
    tokens2,
    attn_weights,
    layer,
    head=head,
    min_line_thickness=0,
    max_line_thickness=5,
    show_all_weights=show_all_weights,
)

Эффективное вычисление весов внимания

Обе картинки в этой лабораторной показывают веса внимания для всех токенов. Как вы увидели, матрицы проекций $Q$ , $K$ и $V$ содержат проекции всех токенов и всех голов. Поэтому ваши вычисления автоматически дают веса внимания и выходы для всех токенов и голов сразу. Это даёт огромные возможности для распараллеливания и делает трансформеры очень удобными для эффективного обучения на больших объёмах данных.

Итоги

В этой лабораторной вы реализовали маскированное многоголовое внимание — центральный механизм любого трансформера. Вы построили маску внимания методом jnp.tri, затем объединили матрицы запросов, ключей и значений всех голов в трёхмерные тензоры. Наконец, вы вычислили веса внимания для всех голов и по ним — выход механизма внимания. Так получаются контекстные эмбеддинги, по которым можно предсказать следующий токен промпта.

Решения

Задание 1

def compute_attention_mask(num_tokens: int) -> jax.Array:
    # Нижнетреугольная матрица из единиц. Форма: (num_tokens, num_tokens).
    attention_mask = jnp.tri(num_tokens)
    return attention_mask

Задание 2

def stack_matrices(
    qkv_dict: dict[str, dict[str, jax.Array]], layer: int
) -> tuple[jax.Array, jax.Array, jax.Array]:
    (
        query_proj_list,
        key_proj_list,
        value_proj_list
    ) = attention.get_qkv_matrices(
        qkv_dict, layer, head=None
    )

    # Новое измерение 0 — голова.
    query_proj = jnp.stack(query_proj_list, axis=0)  # (num_heads, num_tokens, dim_k)
    key_proj = jnp.stack(key_proj_list, axis=0)      # (num_heads, num_tokens, dim_k)
    value_proj = jnp.stack(value_proj_list, axis=0)  # (num_heads, num_tokens, dim_v)

    return query_proj, key_proj, value_proj

Задание 3

def compute_raw_logits(query_proj: jax.Array, key_proj: jax.Array) -> jax.Array:
    # (num_heads, num_tokens, dim_k) -> (num_heads, dim_k, num_tokens).
    key_transposed = jnp.transpose(key_proj, (0, 2, 1))

    _, _, dim_k = key_proj.shape
    # Форма: (num_heads, num_tokens, num_tokens).
    logits = query_proj @ key_transposed / jnp.sqrt(dim_k)

    return logits

Задание 4

def apply_attention_mask(logits_raw: jax.Array) -> jax.Array:
    _, num_tokens, _ = logits_raw.shape

    # Маска (num_tokens, num_tokens) сама «растянется» на все головы.
    attention_mask = compute_attention_mask(num_tokens)
    # Где маска 1 — берём логит, где 0 — K_MASK.
    logits_masked = jnp.where(attention_mask, logits_raw, K_MASK)

    return logits_masked

Задание 5

def compute_attention_weights(logits_masked: jax.Array) -> jax.Array:
    # SoftMax по последнему измерению: в каждой строке веса в сумме дают 1.
    alpha = jax.nn.softmax(logits_masked)
    return alpha

Задание 6

def compute_attention_output(
    alpha: jax.Array, value_proj: jax.Array
) -> jax.Array:
    Y = alpha @ value_proj  # Форма: (num_heads, num_tokens, dim_v).
    return Y

То же самое на PyTorch

Не часть оригинала: всё маскированное многоголовое внимание лабораторной в PyTorch умещается в несколько строк — сравните с уроком «Многоголовое внимание».

import math
import torch

def masked_mha(q, k, v):
    """q, k, v: тензоры формы (головы, токены, d)."""
    n, d_k = q.shape[1], q.shape[2]
    logits = q @ k.transpose(1, 2) / math.sqrt(d_k)          # (головы, n, n)
    mask = torch.tril(torch.ones(n, n, dtype=torch.bool))     # то же, что jnp.tri
    logits = logits.masked_fill(~mask, float("-inf"))         # будущее → −∞
    alpha = torch.softmax(logits, dim=-1)                     # веса внимания
    return alpha @ v, alpha                                   # выход и веса

q = k = v = torch.randn(4, 6, 8)
Y, alpha = masked_mha(q, k, v)
print(Y.shape, alpha[0].sum(dim=-1))   # torch.Size([4, 6, 8]) и строки весов, равные 1

В PyTorch 2 то же делает одна функция torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) — с быстрыми ядрами вроде FlashAttention.

Словарик

English	Русский
attention mechanism, attention weights	механизм внимания, веса внимания
contextual embedding	контекстный эмбеддинг
layer, head	слой, голова (внимания)
beginning of sequence token (BOS)	токен начала последовательности
query, key, value (Q, K, V)	запрос, ключ, значение
projection matrix ( $W_Q$ , $W_K$ , $W_V$ )	матрица проекции
logits	логиты
scaling factor $\sqrt{d_k}$	нормирующий множитель
attention mask, causal mask	маска внимания, причинная маска
masked multi-head attention	маскированное многоголовое внимание
tensor, stack, transpose	тензор, сложить в стопку, транспонировать
reference implementation	эталонная реализация
multi-query attention (MQA), grouped-query attention (GQA)	внимание с общими ключами (MQA), групповое внимание (GQA)
parallelization	распараллеливание

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