Как у нейросетей работает внимание? Статья про self-attention и cross-attention03.10.2024 14:30

Картинка для объяснения процесса

Современные генеративные нейросети, такие как Stable Diffusion или FLUX, создают изображения по текстовым описаниям, используя механизм внимания — attention. Этот механизм помогает моделям как выделять важные части информации, так и связывать промпт с изображением, чтобы в итоге мы получили то, что хотели.

В основе этого процесса лежат два типа внимания: self‑attention (внутреннее внимание), которое определяет взаимосвязи внутри изображения, и cross‑attention (перекрестное внимание), сопоставляющее текстовую подсказку с ее визуальным воплощением.

Например, в случае с запросом «конь на ракете», cross‑attention сопоставит токены («конь» и «ракета») и решит, где разместить коня и ракету, а также в какой обстановке их лучше поместить, например, на Марсе.

Self‑attention, в свою очередь, помогает модели понять, как элементы изображения связаны между собой. Например, как «конь» и «ракета» будут выглядеть вместе и как они состыкуются друг с другом на изображении. Этот процесс обеспечивает логическую связность и целостность финального результата.

Как это работает

Давайте рассмотрим этот процесс на примере: предположим, ваш запрос — конь на ракете с надписью 'быстрый'.

▍Шаг 1: Преобразование текста в эмбеддинги с помощью CLIP

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

▍Шаг 2: Cross-attention — текст направляет генерацию

Cross-attention нужна для связи текстового промпта с объектами на изображении и использует три основные компоненты:

Запросы (Q): векторы, представляющие текущее состояние изображения на шаге генерации (математическое представление изображения).
Ключи (K): векторы, представляющие текстовые эмбеддинги, например, «конь», «ракета», «быстрый».
Значения (V): информация о текстовых эмбеддингах, которая подскажет, как промпт должен повлиять на изображение.

▍Вычисление внимания:

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

$\text{Attention}_{ij} = \frac{Q_i \cdot K_j}{\sqrt{d_k}}$

где:

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

Запросы (Q_i) и ключи (K_j) — это, по сути, векторы, представляющие собой численные описания различных элементов текста и изображения. Запросы исходят от изображения (например, области, где будет конь), а ключи — от текста (например, «ракета», «быстрый», «конь»).

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

Если, например, запрос от изображения на текущем этапе представляет «коня», и ключ — это «ракета», то attention показывает, насколько «ракета» важна для текущего изображения с «конём».

Фичи — это абстрактные представления изображения, которые содержат информацию о его ключевых элементах. Нейросеть использует фичи для того, чтобы научиться распознавать объекты и их взаимосвязи в изображении, а cross-attention выстраивает фичи в единую композицию.

▍Нормализация значений:

Затем Softmax берёт результаты из прошлого шага (то есть числа, показывающие степень важности) и преобразует их в вероятности — это и есть веса $(\alpha_{ij})$ . Чем выше внимание у ключа, тем больший вес у него будет:

$\alpha_{ij} = \text{softmax}(\text{Attention}_{ij})$

Если, например, для «коня» внимание к «ракете» получилось 2, а к «быстрому» — 1.5, то softmax эти значения нормализует, чтобы они сложились в 1 и отразили важность каждого ключа.

▍Взвешивание значений:

После того как softmax вычислил веса, модель использует их для взвешивания значений (V_j) , которые содержат дополнительную информацию о ключах. Например, для «ракеты» это может быть её форма и размер, для «быстрого» — форма надписи на ракете.

Если для «ракеты» вес $(\alpha = 0.6)$ , а для «быстрого» $(\alpha = 0.1)$ то «ракета» окажет большее влияние на финальный результат.

Итоговое значение (Output_i) будет суммой всех этих значений, умноженных на соответствующие веса:

$Output_i = 0.6 \cdot V_{\text{ракета}} + 0.1 \cdot V_{\text{быстрый}} + 0.3 \cdot V_{\text{космос}}$

Это означает, что каждый элемент промпта («ракета», «быстрый») внесёт вклад в результат с учётом своей важности. Если «ракета» важна для коня (вес $(\alpha_{ij})$ для неё больше), то информация о ракете будет больше влиять на текущее изображение.

▍Пример:

Представим, что мы генерируем область с «конём», а ключи — это «ракета», «быстрый», и «космос». В результате внимания мы узнали, что:

Затем применяем softmax для нормализации вероятностей (или веса) $(\alpha_{ij})$ :

«ракета» важна для текущего изображения «коня», и её $(\alpha_{ij}) = 0.6$
«быстрый» тоже важен, но меньше — его $(\alpha_{ij}) = 0.25$
«космос» наименее важен — $(\alpha_{ij}) = 0.15$

Теперь итоговое состояние этого участка изображения будет комбинацией информации о «ракете», «быстром» и «космосе», с учётом их весов:

$Output_i = 0.6 \cdot V_{\text{ракета}} + 0.25 \cdot V_{\text{быстрый}} + 0.15 \cdot V_{\text{космос}}$

В итоге, «ракета» влияет на финальный результат сильнее всего, что скажется на итоговой генерации.

▍Шаг 3: Self-attention — части генерации взаимодействуют друг с другом

После того как cross‑attention отработал и связал текст с изображением, модель использует self‑attention для улучшения связи между различными частями изображения. Например, «конь» и «ракета» должны быть расположены правильно друг относительно друга, а надпись «быстрый» должна находиться на самой ракете.

Self-attention работает по похожему принципу и использует следующие компоненты:

Запросы (Q): векторы, представляющие отдельные смежные пиксели генерации.
Ключи (K) и Значения (V): векторы, представляющие другие части того же изображения (для понимания контекста).

Формула для self-attention:

$\text{Attention}_{ij} = \frac{Q_i \cdot K_j}{\sqrt{d_k}}$

Затем softmax вычисляет веса:

$\alpha_{ij} = \text{softmax}(\text{Attention}_{ij})$

Взвешенные значения влияют на корректировку изображения:

$Output_i = \sum_j \alpha_{ij} V_j$

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

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

▍Шаг 4: Применение CFG Scale

Cross-attention и self-attention работают вместе на каждом шаге, корректируя изображение в соответствии с текстом и внутренними связями изображения.

После работы cross-attention и self-attention модель применяет коэффициент CFG Scale, который позволяет управлять тем, насколько сильно позитивный и негативный промпты влияют на итоговую генерацию.

▍Шаг 5: Sampling method или метод выборки

После того как внимание и CFG Scale откорректировали векторы изображения, модель применяет метод выборки (например, DDIM или другие).

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

Итог

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

Уже только потом подключаются CFG Scale и Sampling method — они являются кистью художника и помогают рисовать там, где нужно, постепенно получая детализированное изображение.

При этом внимание вычисляется на каждом шаге генерации, как CFG Scale и Sampling method.

Заглядывайте в телеграм‑канал, где я пишу гайды по Stable Diffusion и FLUX. Там же будут и анонсы новых статей.