ML: Attention - Механізм уваги

Вступ

Механізм Attention (увага) наразі зустрічається в найрізноманітніших архітектурах і задачах (переклад, генерація тексту, анотація зображень тощо). У цьому документі ми розглянемо історично перше використання Attention у задачі машинного перекладу (Bahdanau D., et al., 2014) і реалізацію функції багатоголової уваги на PyTorch. Наступні документи присвячені застосуванню цього механізму в мережі Transformer і її розвитку в моделях GPT і BERT.

Основна ідея

Перш ніж заглибитися в технічні деталі, розглянемо ідею механізму уваги на прикладі задачі зняття семантичної неоднозначності. При роботі з природною мовою, кожному слову в словнику ставиться у відповідність вектор з дійсними компонентами (ембединг слова). Ці вектори надходять на вхід нейронної мережі з тією чи іншою архітектурою. Компоненти векторів є параметрами, які підбираються в процесі навчання так, щоб близькі за змістом слова мали схожі вектори. Як міра близькості зазвичай виступає косинус кута між векторами, який визначається їхнім скалярним добутком.

Одна з проблем, що виникає при побудові ембедингів слів, - це неоднозначність природної мови (table - це стіл або таблиця тощо). При навчанні на корпусі текстів, вектор ембедингу кожного слова отримує фіксовані компоненти, які не залежать від того сенсу, у якому це слово зустрілося в тексті.

Механізм уваги (точніше в цьому випадку самоуваги) модифікує вектор ембедингу кожного слова, "підмішуючи" до нього вектори його оточення (контексту) з деякими вагами. Нехай, наприклад, у реченні "The table has a lot of data" словам відповідають вектори $\mathbf{v}_1,...,\mathbf{v}_7$. Для модифікації вектора $\mathbf{v}_\text{table}=\mathbf{v}_2$, обчислимо його скалярні добутки з іншими словами: $\{w_1,w_2,...,w_7\}=\{\mathbf{v}_2\mathbf{v}_1,~\mathbf{v}_2\mathbf{v}_2,~...,\mathbf{v}_2\mathbf{v}_7\}$ (включаючи його самого). Отримані ваги віднормуємо за допомогою функції softmax, так щоб їхня сума дорівнювала одиниці, а значення лежали в діапазоні $[0...1]$: $$ w'_i = \text{softmax}(w_1,...,w_7) = e^{w_i}/(e^{w_1}+...+e^{w_7}). $$ Чим ближче до одиниці ваги $w'_i$, тим сильніше слово $\mathbf{v}_2$ схоже на слово $\mathbf{v}_i$. Побудуємо тепер новий ембединг слова table у вигляді зваженої суми: $$ \mathbf{v}'_2 = w'_1\,\mathbf{v}_1+w'_2\,\mathbf{v}_2+...+w'_7\,\mathbf{v}_7. $$ Оскільки артиклі та прийменники у векторному просторі знаходяться далеко від векторів table і data, ми отримаємо щось на кшталт: $$ \begin{array}{llll} \mathbf{v}'_\text{table}\approx 0.9\cdot \mathbf{v}_\text{table}+0.1\cdot \mathbf{v}_\text{data} &~~~~~~~& \text{The table has a lot of data}\\ \mathbf{v}'_\text{table}\approx 0.8\cdot \mathbf{v}_\text{table}+0.2\cdot \mathbf{v}_\text{plate} && \text{There is a plate on the table} \end{array} $$ Для тексту з іншим контекстом (другий приклад вище) вийдуть інші компоненти. У першому випадку "увага" слова table фокусується на слові data, а в другому - на слові plate. У результаті єдиний для всіх смислів вектор $\mathbf{v}_\text{table}$ розщеплюється на два вектори, які мають тенденцію переміщатися до своїх смислових кластерів (row, digits,... у першому випадку і chair, furniture, ... - у другому).

Encoder-Decoder

Нагадаємо, що найпростіша архітектура Encoder-Decoder складається з двох різних рекурентних мереж (енкодера і декодера). Енкодер на вхід отримує текст однією мовою (source), а декодер повинен на виході видати переклад - текст іншою мовою (target). Остання RNN-комірка енкодера на виході містить вектор прихованого стану $\mathbf{h}_e^{\text{last}}$, який "накопичив" у собі інформацію про все source-речення (context vector). Цей вектор відправляють як початковий прихований стан у першу RNN-комірку декодера:

Потім на вхід першої комірки декодера подається службовий токен <BOS> (begin of sentence). На виході цієї комірки мережу навчають видавати слово-переклад "кот" (див. малюнок). Для цього виходи комірок пропускають через лінійний шар (fc) з числом нейронів рівним числу слів у словнику. Потім softmax-функція (sm), видає "ймовірності" слів з яких обирається номер максимальної (argmax). Вектор отриманого слова "кот" передається на вхід другої комірки і т.д. поки не вийде службовий токен <ЕOS> (end of sentence).

Можна також використовувати метод примусового навчання (teacher forcing). У цьому випадку на всі входи декодера одразу подають правильний переклад, а на виході від нього вимагають видати це саме речення зсунуте на одне слово вліво: 

<BOS> кот сидит на коврике   -> [Decoder] ->   кот сидит на коврике <EOS>
Зазвичай між режимами "чесного" і "примусового" навчання проводять випадкове перемикання.

Основна проблема такої архітектури полягає в тому, що при довгих source-реченнях фінальний прихований стан енкодера "забуває" початок речення. Частково цю проблему можна вирішити за допомогою двонапрямних RNN. Однак і в цьому випадку прихований стан, що надходить на вхід декодера, погано "пам'ятає" середину source-речення. Аналогічна проблема виникає в декодері, який з часом "забуває" передане йому source-речення (як прихований стан останньої комірки енкодера). Знизити проблему можна підмішуючи цей прихований стан у приховані стани всіх комірок декодера. Були придумані й інші трюки з боротьби з довгими реченнями. Однак найбільш ефективним виявився механізм уваги.

Фокус уваги в RNN

Зазвичай при перекладі людина декілька разів переглядає source-речення, фокусуючись на важливих на даний момент словах і їхньому оточенні, знімаючи тим самим синтаксичну чи семантичну неоднозначності. Так, слово "mat" може означати "килимок" або "спортивний мат". Для видачі правильного перекладу необхідно сфокусуватися не тільки на слові "mat", а й на слові "cat".

Реалізація цієї ідеї полягає в наступному. До чергового прихованого стану декодера додається зважена сума всіх прихованих станів енкодера. Значення ваг у сумі відображають ступінь важливості того чи іншого слова в source-реченні для генерації поточного слова target-речення.

Будемо позначати вектор прихованого стану комірки енкодера з номером $\alpha$ як $\mathbf{v}_{\alpha}=\{v_{\alpha,0},...,v_{\alpha,E-1}\}$, а декодера як $\mathbf{u}_\alpha=\{u_{\alpha,0},...,u_{\alpha,E-1}\}$, де E - розмірність ембедингу. Ваги уваги комірки декодера з номером $\alpha$ на $\beta$-му слові енкодера позначимо як $w_{\alpha\beta}$. У наступну $(\alpha+1)$-комірку декодера буде відправлятися не поточний прихований стан $\mathbf{u}_\alpha$, а його сума зі зваженими прихованими станами енкодера $\mathbf{u}_\alpha +\mathbf{u}'_\alpha$, де: $$ \mathbf{u}'_\alpha = \sum_\beta w_{\alpha\beta} \,\mathbf{v}_\beta,~~~~~~~~~~~\sum_\beta w_{\alpha\beta} = 1,~~~~~~~~~~~ w_{\alpha\beta} = f\bigr(\mathbf{u}_\alpha, \mathbf{v}_\beta \bigr). $$ Як функція $f$, що обчислює ваги уваги $w_{\alpha\beta}$, може виступати повнозв'язний шар або скалярний добуток $\mathbf{u}_\alpha\cdot \mathbf{v}_\beta$ (вектори близьких слів паралельні). Нижче на малюнку функція $f$ зображена синім кругом. Сума ваг $w_{\alpha\beta}$ за другим індексом $\beta$ повинна дорівнювати одиниці, що забезпечує шар softmax. Приховані стани $v_\alpha$, як зазвичай, потрапляють як на вихід (у функцію $f$), так і в наступну комірку рекурентного шару:

У даному прикладі вектор прихованого стану, отриманий при подачі службового слова "<BOS>" (begin of sentence) і останнього прихованого стану енкодера, повинен найближче виявитися до векторів перших комірок енкодера (для однотипних за синтаксисом мов). До нього також підмішується вектор слова sits і з меншими вагами наступні слова вихідного речення. Друга комірка декодера, отримавши на вхід вектор слова "кот", повинна сфокусуватися на слові "cat" і семантично близьких словах ("sits" тощо).

Підкреслимо, що source і target мови мають різні ембединги. Однак у процесі навчання вектори слів "cat" і кот виявляються близькими, що приводить до збільшення відповідних ваг уваги.

Якщо функція $f$ пропорційна скалярному добутку векторів ($\mu=\text{const}$): $$ f(\mathbf{u},\mathbf{v}) = \mu\,\mathbf{u}\cdot\mathbf{v}, $$ то ваги уваги до прихованих станів декодера визначаються таким чином: $$ w_{\alpha\beta} = \text{softmax}\bigr(\mu\,\mathbf{u}_\alpha\mathbf{v}_\beta\bigr) = \frac{e^{\mu\,\mathbf{u}_\alpha\mathbf{v}_\beta} } {\sum_\gamma e^{\mu\, \mathbf{u}_\alpha\mathbf{v}_\gamma } }. $$

Ще один варіант функції обчислення ваг містить дві матриці $\mathbf{W}_1$, $\mathbf{W}_2$ і вектор $\mathbf{b}$ (компоненти яких підбирають у процесі навчання): $$ f(\mathbf{u},\mathbf{v}) = \tanh\bigr(\mathbf{u}\cdot\mathbf{W}_1+\mathbf{v}\cdot\mathbf{W}_2\bigr)\,\mathbf{b}. $$

Як ваги далі будемо використовувати скалярний добуток $\omega_{\alpha\beta} = \mu\,\mathbf{u}_\alpha\cdot\mathbf{v}_\beta$ і запишемо результат дії механізму уваги в матричному вигляді.

Функція Attention

Нехай є три матриці: матриця $\mathbf{Q}$ запитів (query), матриця $\mathbf{K}$ ключів (key) і матриця $\mathbf{V}$ значень (value). Введені матриці є аргументами функції уваги: $$ \mathbf{A} = \text{Attn}( \mathbf{Q},\, \mathbf{K},\, \mathbf{V}) ~=~ \text{softmax} \Bigr(\frac{\mathbf{Q}\cdot\mathbf{K}^\top}{\sqrt{E}}\Bigr)\,\mathbf{V}, $$ де $\top$ - операція транспонування, що переставляє місцями індекси матриці і функція softmax незалежно застосовується до кожного рядка матриці $\mathbf{Q}\cdot\mathbf{K}^\top/\sqrt{E}$. У загальному випадку матриці можуть мати такі форми: $$ \mathbf{Q}:~~~(N,\,E),~~~~~~~~~~~\mathbf{K}:~~~(M,\,E),~~~~~~~~~~~\mathbf{V}:~~~(M,\,E'),~~~~~~~~~~~\mathbf{A}:~~~(N,\,E'). $$

Назви введених матриць пов'язані з простою моделлю асоціативної пам'яті. Будемо вважати, що в пам'яті зберігається $M$ пар векторів ключ-значення: $(\mathbf{k}_1,\mathbf{v}_1)$,...., $(\mathbf{k}_M,\mathbf{v}_M)$. Вони є рядками матриць $\mathbf{K}$ і $\mathbf{V}$, маючи, взагалі кажучи, різні розмірності $E$ і $E'$. Нехай матриця запиту $\mathbf{Q}$ складається з одного рядка ($N=1$) - вектора $\mathbf{q}$ розмірності $E=2$ і є $M=3$ ключа $k_{\alpha i}$ (перший індекс - номер ключа, другий - номер його компоненти). Тоді: $$ \mathbf{Q}\,\mathbf{K}^\top ~=~ \begin{array}{|c|c|} \hline q_0 & q_1 \\ \hline \end{array} \cdot \begin{array}{|c|c|} \hline k_{00} & k_{00} \\ \hline k_{10} & k_{10} \\ \hline k_{20} & k_{20} \\ \hline \end{array}^{~\top} ~=~ \begin{array}{|c|c|} \hline q_0 & q_1 \\ \hline \end{array} \cdot \begin{array}{|c|c|c|} \hline k_{00} & k_{10} & k_{20}\\ \hline k_{01} & k_{11} & k_{21}\\ \hline \end{array} ~=~ \begin{array}{|c|c|} \hline \mathbf{q}\mathbf{k}_0 & \mathbf{q}\mathbf{k}_1 & \mathbf{q}\mathbf{k}_2\\ \hline \end{array} $$ Якщо вектор запиту $\mathbf{q}$ паралельний $i$-тому ключу $\mathbf{k}_i$ і антипаралельний решті ключів $\mathbf{q}\mathbf{k}_j \ll -1$, $j\neq i$, то софтмакс поверне вектор з нулів, крім $i$-ї позиції, де буде стояти 1. Відповідно, функція Attn поверне значення $\mathbf{v}_i$ з $i$-тої пари. У загальному випадку, якщо точного збігу ключа і запиту немає, то повертається зважена сума з переважанням значень, ключі яких найбільш схожі (со-спрямовані) запиту.

Відповідно до форм матриць $\mathbf{Q},\,\mathbf{K}$, аргументом функції softmax виступає матриця форми $(N,M)$, оскільки $(N,E) \cdot (M,E)^\top= (N,E) \cdot (E,M) = (N,M)$. Функція softmax, застосована до кожного рядочка цієї матриці (за розмірністю $M$), дає матрицю $\Omega_{ij}$ тієї самої форми. Сума її елементів за другим індексом $j$ дорівнює одиниці для кожного $i$. Ця матриця згортається з $\mathbf{V}$, тому результатом дії функції буде матриця форми $(N,\,E')$. Якщо $N > 1$, то функція $\text{Attn}$ одночасно обробляє декілька запитів (рядків матриці $\mathbf{Q}$).

☝ Матриця, обернена до матриці $\mathbf{A}$, позначається як $\mathbf{A}^{-1}$, Матричний добуток $\mathbf{Q}\cdot\mathbf{K}^\top$ запитів і ключів, слідуючи Vaswani A., et al. (2017), прийнято ділити на корінь з розмірності $E$ векторів матриць $\mathbf{Q}$ і $\mathbf{K}$. Мотивація для цього може бути такою. Нехай компоненти двох векторів $\mathbf{q}$ і $\mathbf{k}$ є незалежними випадковими величинами з нульовим середнім і одиничною дисперсією. Тоді скалярний добуток $\mathbf{q}\cdot\mathbf{k}$ також має нульове середнє і дисперсію рівну розмірності векторів $E$, тобто типові значення $\mathbf{q}\cdot\mathbf{k}$ знаходяться в інтервалі $\pm\sqrt{E}$. Масштабування переводить їх до інтервалу $\pm 1$.

Повернемося тепер до механізму уваги, описаного в попередньому розділі. Нехай $\mathbf{U}= u_{\alpha i}$ - вектори прихованих станів декодера, а $\mathbf{V}= v_{\alpha i}$ - енкодера, де перший індекс - номер вектора, а другий - його компоненти (кожен рядок матриць $\mathbf{U}$ і $\mathbf{V}$ це $\alpha$-й вектор). Тоді, компоненти векторів $\mathbf{u}'_\alpha$ - добавок до прихованого стану декодера є рядками матриці ($\mu = 1/\sqrt{E}$): $$ \mathbf{U}' = \text{Attn}( \mathbf{U},\, \mathbf{V}, \, \mathbf{V}). $$ У цьому випадку $M$ - число комірок енкодера ( = числу слів вхідного речення = пар ключ-значення), а $N=1$ (одне слово поточної комірки декодера).

Реалізація функції Attn на PyTorch має такий вигляд (B - число прикладів у батчі, E - розмірність ембедингу, N - число запитів; M число пар ключ-значення): 
import math
import torch
import torch.nn as nn

def Attn(Q, K, V):                                        # Q: (B,N,E); K,V: (B,M,E)
    E = Q.size(-1)                                        # розмірність ембедингу E
    W = torch.bmm(Q, K.transpose(-2, -1)) / math.sqrt(E)  # (B,N,M)        
    W = nn.functional.softmax(W, dim = -1)                # за останнім індексом тензора
    return torch.bmm(W, V)                                # (B,N,E)

Метод bmm перемножує дві матриці: bmm( (B,N,E), (B,E,M) ) = (B,N,M) незалежно для кожного прикладу в батчі, тобто у "циклі" по прикладах B звичайним чином множить: (N,E) @ (E,M) = (N,M).

Самоувага

Механізм уваги можна застосовувати не тільки в архітектурі енкодер-декодер. Розглянемо, наприклад, проблему неоднозначності сенсу слів, згадану на початку документа. Нехай є проста архітектура мережі з $N$ входами і $N$ виходами. На її входи надходять $N$ слів (у вигляді $E$-вимірних векторів ембедингу), а на виходах повинні вийти змінені вектори цих же слів, з урахуванням контексту всього речення.

Нехай матриця $\mathbf{V}$ розмірності $(N,E)$ у $N$ рядках містить $N$ векторів слів речення. Обчислимо значення функції $\text{Attn}(\mathbf{V},\, \mathbf{V},\,\mathbf{V})$. Оскільки матриця запиту збігається з матрицями ключів і значень, ця ситуація називається само-увагою (self-attention). Розпишемо матричні множення в явному вигляді, наприклад, для $N=3$ і $E=2$ (у матриці $\mathbf{V}:~v_{\alpha i}$, де, як і раніше, перший індекс - це номер слова, другий - номер компоненти його вектора): $$ \mathbf{V}\cdot\mathbf{V}^\top ~=~ \begin{array}{|c|c|} \hline v_{00} & v_{01} \\ \hline v_{10} & v_{11} \\ \hline v_{20} & v_{21} \\ \hline \end{array} \cdot \begin{array}{|c|c|} \hline v_{00} & v_{10} & v_{20} \\ \hline v_{01} & v_{11} & v_{21} \\ \hline \end{array} = \begin{array}{|c|c|c|} \hline \mathbf{v}_0\mathbf{v}_0 & \mathbf{v}_0\mathbf{v}_1 & \mathbf{v}_0\mathbf{v}_2 \\ \hline \mathbf{v}_1\mathbf{v}_0 & \mathbf{v}_1\mathbf{v}_1 & \mathbf{v}_1\mathbf{v}_2 \\ \hline \mathbf{v}_2\mathbf{v}_0 & \mathbf{v}_2\mathbf{v}_1 & \mathbf{v}_2\mathbf{v}_2 \\ \hline \end{array} $$ У першому рядку результуючої матриці знаходяться скалярні добутки вектора $\mathbf{v}_0$ першого слова з усіма іншими словами речення. У другому - аналогічні добутки другого слова $\mathbf{v}_1$ і т.д. Найбільші (додатні) значення зазвичай мають діагональні елементи матриці $\mathbf{v}^2_0$, $\mathbf{v}^2_1$, $\mathbf{v}^2_2$ (при порівнянних довжинах векторів). Після обчислення (порядкового) функції софтмаксу вийде щось на кшталт: $$ \text{Attn}(\mathbf{V},\, \mathbf{V},\,\mathbf{V}) ~=~ \text{softmax}\Bigr(\frac{\mathbf{V}\cdot\mathbf{V}^\top}{\sqrt{2}}\Bigr)\cdot\mathbf{V} ~=~ \begin{array}{|c|c|c|} \hline 0.8 & 0.2 & 0 \\ \hline 0.3 & 0.7 & 0 \\ \hline 0.1 & 0 & 0.9 \\ \hline \end{array} \cdot \begin{array}{|c|c|} \hline v_{00} & v_{01} \\ \hline v_{10} & v_{11} \\ \hline v_{20} & v_{21} \\ \hline \end{array} ~=~ \begin{array}{|l|l|} \hline 0.8\,v_{00} + 0.2\,v_{10} & 0.8\,v_{01} + 0.2\,v_{11} \\ \hline 0.7\,v_{10} + 0.3\,v_{00} & 0.7\,v_{11} + 0.3\,v_{01} \\ \hline 0.9\,v_{20} + 0.1\,v_{00} & 0.9\,v_{21} + 0.1\,v_{01} \\ \hline \end{array} $$ Таким чином, для трьох слів на виході отримуємо три вихідних вектори до яких "підмішані" компоненти близьких за змістом слів оточення (якщо ембединг був вдало побудований).

Багатоголовий Attention

Наступним етапом розвитку технології Attention, стало введення декількох навчуваних фокусів уваги на "різних аспектах" послідовності ключів. Нагадаємо, що в згорткових мережах число вторинних ознак збільшують, додаючи на черговому рівні аналізу зображення додаткові фільтри. Аналогічно, при обробці послідовностей, вводять H фокусів уваги, які називають головами (кожна голова "дивиться у свій бік"): $$ \mathbf{A}~=~\text{MultiHead}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{Concat}(\mathbf{h}_1,\dots,\mathbf{h}_H)\,\mathbf{W}^O, ~~~~~~~~~~~\mathbf{h}_i = \text{Attn}(\mathbf{Q}\,\mathbf{W}_i^Q,~ \mathbf{K}\,\mathbf{W}_i^K,~ \mathbf{V}\,\mathbf{W}_i^V). $$

Спочатку для векторів запиту, ключів і значень роблять лінійне перетворення за допомогою трьох матриць $\mathbf{W}^Q_i$, $\mathbf{W}^K_i$, $\mathbf{W}^V_i$. При цьому для кожної ($i$-тої) голови $\mathbf{h}_i$ набір матриць свій (свій "фокус уваги"). Якщо: $$ \mathbf{Q}:~(N,E),~~~~~~~\mathbf{K}:~(M,E_k),~~~~~~~\mathbf{V}:~(M,E_v),~~~~~~~~~~~~\mathbf{A}:~(N,E_a), $$ то, у загальному випадку, матриці можуть мати такі форми: $$ \mathbf{W}^Q_i:~(E,\,E_i),~~~~~~~~~~\mathbf{W}^K_i:~(E_k,\,E_i),~~~~~~~~~~\mathbf{W}^V_i:~(E_v,\,E_h),~~~~~~~~\mathbf{h}_i:~(N,\,E_h),~~~~~~~ \mathbf{W}^O:~(E_h\, H,~E_a). $$ Після обчислення функції уваги Attn для кожної голови, виходять матриці форми $(N,E_h)$. Їх об'єднують в одну (конкатенують) за останнім індексом, що дає матрицю $(N,~E_h\, H)$. Її згортка з $\mathbf{W}^O$ приводить до фінальної матриці $\mathbf{A}$ форми $(N,\,E_a)$.

У вихідній статті (2017), де були введені голови уваги, було покладено $E_i=E_h=E/H$. Ця сама угода прийнята в PyTorch і матриці всіх голів упаковуються в одну з E колонками. Для задачі само-уваги ($E_k=E_v=E$) така матриця має квадратну форму $(E,E)$. $$ \mathbf{W}^Q_i:~(E,\, E/H),~~~~~~~~\mathbf{W}^K:~(E_k,\,E/H),~~~~~~~~~~~~~\mathbf{W}^V_i:~(E_v,\,E/H),~~~~~~~~~~~~\mathbf{W}^O:~(E, ~E). $$ При такому виборі, чим більше голів, тим більше різних аспектів послідовності вони можуть "розгледіти". Однак кожна голова при цьому "бачить гірше", оскільки оперує ембедингом розмірності $E/H$.

За фактом вхідні матриці $\mathbf{Q},\,\mathbf{K},\,\mathbf{V}$ спочатку піддаються лінійному перетворенню, а потім розрізаються по вертикалі на $H$ голів і ваги уваги обчислюються вже незалежним чином для кожної голови.

🔥 Зауважимо, що у функції уваги як параметри виступають елементи матриці добутку $\mathbf{W}^Q\cdot\mathbf{W}^{K\top}:~(E,E_k)$, а не ці дві матриці окремо. Однак, якщо $E_k=E$ і $E_i=E/H$, число елементів двох матриць дорівнює $2\,E^2/H$, що при трьох і більше головах менше $E^2$ елементів їхнього добутку.

MultiheadAttention в PyTorch

У PyTorch є готова функція багатоголової уваги, оформлена як nn-шар:

✒ nn.MultiheadAttention
… (embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False,
… add_zero_attn=False, kdim=None, vdim=None)

Аргументи конструктора мають такі значення: embed_dim = E - розмірність ембедингу; num_heads = H - число голів; $\text{kdim}=E_k$, $\text{vdim}=E_v$ - розмірності ключів і значень (якщо вони None, то вважаються рівними E). Розмірність ембедингу E повинна бути кратна числу голів H. При заданні параметрів bias, add_bias_kv у True проводиться не тільки множення запитів, ключів і значень на матриці, а й зсув на вектор (загальне лінійне перетворення). Пряме поширення через модуль уваги має такі параметри:

✒ nn.MultiheadAttention.forward
… (query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None)

У найпростішому випадку на вхід об'єкта MultiheadAttention подаються тензори: $\mathbf{Q},~\mathbf{K},~\mathbf{V}$, а на виході виходить пара ($\mathbf{A}$, $\mathbf{W}$) - результат уваги $\mathbf{A}$ і тензор ваг $\mathbf{W}$ (дія функції softmax для кожного прикладу батчу):

де N - довжина послідовності запиту (для декодера довжина target-речення), M - число ключів і значень (довжина послідовності source для енкодера), B - розмір батчу (число оброблюваних прикладів). Якщо ваги не потрібні, можна покласти need_weights=False (але функція уваги все одно буде повертати пару, у якій другий елемент дорівнює None).

Зверніть увагу, що за аналогією з рекурентними шарами індекс номера прикладу в батчі B йде не першим, а другим (крім результуючого тензора ваг). Наведемо приклад використання багатоголової уваги: 

E, H, N, M, B = 100, 10, 3, 3, 1
Ek, Ev        = 100, 100

Q, K, V = torch.rand(N, B, E), torch.rand(M, B, Ek), torch.rand(M, B, Ev) 

MHA = nn.MultiheadAttention(E, H, kdim=Ek, vdim=Ev)
A,W = MHA(Q, K, V)

print(tuple(A.shape), tuple(W.shape))               # (3, 1, 100) (1, 3, 3)

Розмірності матриць дорівнюють: $\mathbf{W}^Q:$ (E, E), $\mathbf{W}^K:$ (E, kdim), $\mathbf{W}^V:$ (E, vdim), де розмірності kdim,vdim або задаються, або рівні E (вони передаються у функцію linear, тому при множенні транспонуються). Розмірність кожної голови дорівнює E // H. Вихідна матриця $\mathbf{W}^O:$ це шар Linear(E, E, bias=bias).

Якщо ембединг входів однаковий, то проєкційні матриці упаковані в in_proj_weight форми (3*E,E), інакше це три різні матриці: q_proj_weight, k_proj_weight, v_proj_weight.
Зміщення (якщо воно є) це in_proj_bias форми (3*E,), а вихідний шар: out_proj: 

for k, v in MHA.state_dict().items():               # in_proj_weight   shape: (300, 100) 
    print(f'{k:20s}  shape: {tuple(v.shape)} ')     # in_proj_bias     shape: (300,) 
                                                    # out_proj.weight  shape: (100, 100) 
                                                    # out_proj.bias    shape: (100,)

Маскована увага

Важливими параметрами прямого поширення через шар багатоголової уваги є маски key_padding_mask і attn_mask.

Булева маска key_padding_mask: (B,M) дозволяє виключати з механізму уваги деякі пари ключів і значень (незалежно для кожного прикладу батчу B). Для цього номери виключених пар у масці необхідно позначити значенням True.

Дійсна маска attn_mask: (N,M) додається до ваг уваги перед їхньою нормуванням за допомогою функції софтмакс. Зазвичай ця маска використовується для точкового відключення конкретних ваг. Для цього відповідні елементи маски покладаються рівними мінус нескінченності, а решта - нулю: 

E, H, N, M, B = 8, 4, 3, 4, 1
Q, K, V = torch.rand(N,B,E), torch.rand(M,B,E), torch.rand(M,B,E) 

MHA = nn.MultiheadAttention(E,H)

A,W = MHA(Q, K, V, 
          key_padding_mask=torch.tensor( [[False,True,False,True]]) )
print(W)

inf = float("-inf")
A,W = MHA(Q, K, V, 
          attn_mask=torch.tensor( [[inf, 0.0, 0.0, inf],
                                   [0.0, 0.0, 0.0, 0.0],
                                   [inf, 0.0, 0.0, inf]] ) )
print(W)
У першому випадку key_padding_mask відключила ключі з індексами 1 і 3 Нижче нулями забиті відповідні їм колонки матриці ваг W. Маска attn_mask "видалила" кутові елементи матриці ваг уваги: 
[[[0.5148, 0.0000, 0.4852, 0.0000],              [[[0.0000, 0.4940, 0.5060, 0.0000],
  [0.5136, 0.0000, 0.4864, 0.0000],                [0.2570, 0.2519, 0.2435, 0.2475],
  [0.5127, 0.0000, 0.4873, 0.0000]]]               [0.0000, 0.4918, 0.5082, 0.0000]]],
З документа, присвяченого трансформеру, стане зрозумілою практична користь від застосування цих масок.

Реалізація MultiheadAttention

Відтворимо обчислення, що відбуваються всередині шару nn.MultiheadAttention. Нам знадобляться такі функції PyTorch: 

from torch               import bmm                   # пакетне множення матриць 
from torch.nn.functional import linear  as linear     # лінійна функція y = x@A^T + b
from torch.nn.functional import softmax as softmax    # функція софтмаксу
У функцію MultiHeadAttention, крім матриць запитів Q, ключів K, значень V і масок key_mask, attn_mask, передамо також матриці лінійних перетворень, які візьмемо у шару nn.MultiheadAttention. У коментарях, як зазвичай, наведено форми тензорів, що виходять: 
def MultiHeadAttention(Q,K,V,                         # Q:(N,B,E);  K:(M,B,Ek);  V:(M,B,Ev)
                       Wq, Wk, Wv, Wo,                # матриці повороту  для  Q,K,V,A
                       Bq=None, Bk=None, Bv=None, Bo=None,  # матриці зміщення для  Q,K,V,A
                       key_mask  = None,              # маска виключення ключів (B,M)
                       attn_mask = None):             # адитивна маска        (N,M)
    
    q = linear(Q, Wq, Bq)                             # (N,B,E) лінійне перетворення
    k = linear(K, Wk, Bk)                             # (M,B,E)
    v = linear(V, Wv, Bv)                             # (M,B,E)
    
    q = q.view(N, B*H, E//H).transpose(0,1)           # (B*H, N, E/H)  розбиваємо на H голів
    k = k.view(M, B*H, E//H).transpose(0,1)           # (B*H, M, E/H)
    v = v.view(M, B*H, E//H).transpose(0,1)           # (B*H, M, E/H)
    
    W = bmm(q, k.transpose(1,2))*float(E//H)**-0.5    # (B*H, N, M) розмірність голови E/H
    
    if attn_mask is not None:                         # адитивна маска
        W += attn_mask.unsqueeze(0)                   # (N,M) -> (1,N,M)
 
    if key_mask is not None:                          #  виключаємо частину ключів
        W = W.view(B, H, N, M)
        key_mask = key_mask.unsqueeze(1).unsqueeze(2) # (B,1,1,M)
        W = W.masked_fill(key_mask, float('-inf'))
        W = W.view(B*H, N, M)
    
    W = softmax(W, dim=-1)                            # (B*H, N, M)

    A = bmm(W, v)                                     # (B*H, N,  E/H)

    A = A.transpose(0, 1).contiguous().view(N,B,E)    # (N, B, E)
    W = W.view(B, H, N, M)
    return linear(A, Wo, Bo), W.sum(dim=1) / H        # (N, B, E), (B, N, M)

Спочатку Q, K, V піддаються лінійному перетворенню linear. Потім вони розрізаються на голови. Для цього змінюється їхня форма і перші два індекси переставляються місцями. У результаті перший індекс нумерує і приклади, і голови. Використання далі методу множення bmm проводить перемноження прямокутних матриць для кожного прикладу, кожної голови. Транспонування і зміна форми тензора A в кінці функції переводить його у вихідну форму, подібну до тензора Q.

У блоках if проводиться маскування уваги. Спочатку до ваг W додається маска attn_mask до якої додається одна розмірність (і далі вмикається механізм розширення). Потім булева маска key_mask слугує для заміни колонок для яких key_mask == True на мінус нескінченність. Після проходження через функцію софтмакс, елементи матриці ваг у цих колонках дорівнюватимуть нулю.

Якщо об'єкт MHA створено за допомогою конструктора nn.MultiheadAttention (як у попередньому розділі), то виклик написаної функції, матиме такий вигляд: 
A2,W2 = MultiHeadAttention(Q,K,V, 
                           MHA.in_proj_weight[0 : E]   if Ek==E else  MHA.q_proj_weight, 
                           MHA.in_proj_weight[E : 2*E] if Ek==E else  MHA.k_proj_weight, 
                           MHA.in_proj_weight[E*2 : ]  if Ek==E else  MHA.v_proj_weight, 
                           MHA.out_proj.weight,                           
                           MHA.in_proj_bias[0 : E], 
                           MHA.in_proj_bias[E : E*2], 
                           MHA.in_proj_bias[E*2:], 
                           MHA.out_proj.bias,
                           attn_mask = attn_mask,
                           key_mask = key_mask  )

print(((A1.detach() - A2.detach())**2).sum()**0.5)   # порівнюємо матриці (округлення)
print(((W1.detach() - W2.detach())**2).sum()**0.5)

Література

Статті

Вихідні матеріали