ML: Embedding слів

Шар Embedding

Методи ембедингу, засновані на контекстних частотах і PCA-зниженні розмірності, важко використовувати при словниках у кілька сотень тисяч слів. Крім цього вони не дозволяють покращувати якість векторизації шляхом ускладнення моделі. На допомогу, як зазвичай, приходять нейронні мережі.

У фреймворках з роботи з нейронними мережами існує спеціальний шар Embedding. Цей шар на вхід отримує номери слів, а на виході видає їхні векторні подання (до початку навчання вони випадкові):

Вище E = 2 і у шару три входи (nX = 3). У першого слова номер 0, у другого 2, а у третього 1. Шар Embedding зберігає матрицю форми (V, E), з якої, при подачі на вхід числа i, видає i-й рядок.

Векторизацію (= ембединг = Word2Vec) слів словника отримують різними способами. Наприклад, можна одразу навчати матрицю шару Embedding безпосередньо на своїй задачі. Можна також попередньо отримати "грубі" значення компонент векторів на спрощеній моделі, а потім дообучати на складнішій задачі. Зазвичай, терміни векторизація (або ембединг) використовують у загальному випадку. Термін Word2Vec вужчий і частіше відноситься до описуваних нижче методів Skip-gram і CBOW.

Skip-gram і CBOW

Нехай є довга послідовність слів (наприклад, природної мови): $w_0,w_1,...$. Слова близькі за змістом повинні знаходитися в схожому оточенні і вектори таких слів повинні бути близькими.

У Skip-gram як навчальні дані використовуються пари: дане слово $w_t$ і одне слово з його оточення. Для складання таких впорядкованих пар формуються $(n=2m+1)$-грами, де $m$ дорівнює числу слів до $w_t$ і після нього: $$ \{w_{t-m},...,w_{t-1},\,\underline{w_{t}}\,,w_{t+1},...,w_{t+m}\} ~~~~~~\Rightarrow~~~~~~ \{~ (w_t, w_{t\pm i}),~~~~~i=1...m~\}. $$

Цей метод існує у двох варіантах. У варіанті Skip-gram Softmax за словом $\mathbf{X} = w_t$ передбачають одне зі слів оточення $w_{t\pm i}$. На виході мережі з V нейронами розташовують softmax шар, що дає "умовні ймовірності": $\mathbf{Y} ~=~ p(w_i| w_t),~~~i=0...V-1$ усіх слів у словнику. Функція $\text{argmax}_{i}\,p(w_i|w_t)$ дає номер $i$ найбільш ймовірного слова.

У Skip-gram Negative Sampling обидва слова пари $\mathbf{X}=(w_t, w_{t\pm i})$ надходять на вхід і відносяться до позитивного класу ($Y=1$). Пари $X=(w_t, w')$, де $w'$ - випадкове слово тексту, відносять до негативного класу ($Y=0$). Якщо словник досить великий, то випадкове слово $w'$ швидше за все виявиться не з оточення $w_t$. Вибір випадкового слова з тексту, а не словника імітує його правильну частотність. На виході мережі знаходиться єдиний нейрон із сигмоїдною функцією = [0...1] (бінарна класифікація).

CBOW (Continuous Bag of Words) - альтернативний метод Word2Vec, у якому за всіма векторами слів, що оточують слово $w_t$, передбачається його "ймовірність" (на виході розмірності V знаходиться функція softmax): $$ \mathbf{X} = [w_{t-m},...,w_{t-1},w_{t+1},...,w_{t+m}],~~~~~~~~~~Y = w_{t}. $$

Нижче ми розглянемо приклади архітектур цих методів. Вважається, що Skip-gram краще працює для рідкісних слів (оскільки вони не "ховаються" при усередненні SBOW). Хоча при цьому SBOW краще використовує контекст для передбачення слова.

Як і у випадку з ймовірнісним методом, розглянутим вище, ці методи групують разом семантично схожі поняття. Розглянемо, наприклад, назви днів тижня ("Monday", "Tuesday",..., "Sunday"). Вони рідко зустрічаються поруч в одному реченні, однак знаходяться у схожих контекстах: "on Monday morning", "every Monday", "It was a Monday", "one Monday", де на місці Monday може знаходитися будь-який день тижня. У результаті, усі дні тижня, повертаючи свої вектори в бік векторів контексту, прагнуть притягнутися до загальної області векторного простору, куди "веде сумарна сила" від контекстних слів.

Пари слів для Skip-gram методу

Пари pairs для методу Skip-gram будемо складати зі слів, що потрапили в симетричне вікно WIN навколо центрального слова в списку words усередині речення: 
def get_pairs(words, WIN, pairs):
    for i1, w1 in enumerate(words):               
        if not w1 in word_to_id: continue
    
        i2_beg = max(0, i1-WIN)
        for i2 in range(i2_beg, min(len(words), i1+WIN+1) ):  # навколо  w1            
            if i2 != i1 and words[i2] in word_to_id:
               pairs.append((w1, words[i2]))
Наприклад після речення "['the', 'cat', 'likes', 'the', 'mat']" для вікна WIN=2 вийде такий набір пар: 
('the',   'cat'), ('the',   'likes'), 
('cat',   'the'), ('cat',   'likes'), ('cat',   'the'), 
('likes', 'the'), ('likes', 'cat'),   ('likes', 'the'), ('likes', 'mat'), ...

Оскільки речення у використовуваному корпусі досить короткі, ми візьмемо вікно WIN великим, вважаючи, що всі слова речення утворюють контекст для будь-якого його слова: 

pairs = []                                     # список пар
for doc in docs:                               # по документах
    for sent in doc:                           # по реченнях документа
        get_pairs(sent.split(), 100,  pairs)

Перед побудовою навчальних прикладів, як зазвичай, пари варто перемішати: 
import random
random.shuffle(pairs)                          # випадково перемішуємо
У корпусі ROCStories (див. файл 100KStories.zip), при складанні пар усередині речення для WIN = 100, виходить 45'671'880 пар (при словнику в 10'393 слова).

Skip-gram Softmax

У парі $(\mathbf{u},\mathbf{w}_i)$ вектор "центрального" слова $\mathbf{u}$ і вектор контексту $\mathbf{w}_i$ будемо вважати такими, що належать до різних ембедингів (тобто відрізняємо "центральне" слово від слів з його оточення). Це означає, що $\mathbf{u}$ і $\mathbf{w}$ є рядками різних матриць однакової форми (V, E). Ступінь близькості слів характеризується скалярним добутком векторів $\mathbf{u}\mathbf{w}_i$ (паралельні вектори семантично близькі). Отримуючи на вхід моделі вектор центрального слова $\mathbf{u}$, на виході будемо передбачати розподіл ймовірностей слів його оточення. Для цього використовується функція softmax (у знаменнику сума по всіх словах словника $\mathbf{w}_j$): $$ p_i = \frac{e^{\mathbf{u}\mathbf{w}_i}}{\sum_j e^{\mathbf{u}\mathbf{w}_j}}. $$ Якщо пара складається зі слів з номерами (i,j), то на вхід $\mathbf{X}$ моделі будемо подавати ціле число i (номер слова), а на виході $\mathbf{Y}$ очікувати вектор p=[0,0,...,0,1,0,..,0], де 1 стоїть на j-му місці. З метою економії пам'яті, для виходу $Y$ вказується тільки номер j слова (sparse-кодування для CrossEntropyLoss у PyTorch): 

X = np.zeros((len(pairs), 1),   dtype=np.int64)
Y = np.zeros((len(pairs)),      dtype=np.int64)    

for i, (u,w) in enumerate(pairs):        
    X[i,0] = wordID[u]["id"]          # індекс вхідного (центрального) слова
    Y[i]   = wordID[w]["id"]          # індекс передбачуваного слова

Найпростіша реалізація Skip-gram методу в PyTorch має вигляд: 

E_DIM   =  100                                               # розмірність векторів

model = nn.Sequential(                                       # (N,1)   input X
        nn.Embedding(V_DIM, E_DIM, scale_grad_by_freq=True), # (N,1,E)
        nn.Flatten(),                                        # (N,E)
        nn.Linear   (E_DIM, V_DIM, bias=False))              # (N,V)   CrossEntropyLoss

Шар Embedding зберігає вектори "центрального" слова і має V*E параметрів.
Вхід $\mathbf{X}$ має форму (N,1), де N - число прикладів у батчі. Після Embedding вийде тензор $\mathbf{X}_e:$ (N,1,E). Шар Flatten змінює його розмірність на $\mathbf{X}_f:$ (N,E). Параметр scale_grad_by_freq повідомляє, що при навчанні вектори рідкісних слів у батчі треба зсувати сильніше, ніж часто вживаних слів. Природно, при цьому розмір батчу варто збільшити (наприклад N=512).

Повнозв'язний шар Linear (без зміщення bias) містить матрицю $\mathbf{W}$ ваг форми (V, E) з такою самою кількістю параметрів, що і в Embedding. Ця матриця є списком векторів $\mathbf{w}_i$. У лінійному шарі відбувається множення $\mathbf{X}_f\cdot\mathbf{W}^\top,$ де $\mathbf{W}^\top$ - транспонована матриця. Оскільки в кожному рядку навчальних даних знаходиться вектор вхідного слова $\mathbf{u}$, він згортається з векторами $\mathbf{w}_i$ і після проходження через softmax виходять ймовірності $p_i$. У PyTorch функція softmax обчислюється всередині помилки CrossEntropyLoss, тому на виході мережі її ставити не треба.

$$ \text{N} \left\{ \phantom{ \begin{array}{} \\ \\ \\ \\ \end{array} } \right. \overbrace{ \underbrace{ \begin{array}{|c|c|c|} \hline u_1 & u_2 & u_3 \\ \hline ~ & ~ & ~ \\ \hline ~ & ~ & ~ \\ \hline ~ & ~ & ~ \\ \hline \end{array} }_ {\mathbf{u}} }^ {\displaystyle\mathrm{E}} ~~~ \cdot ~~~ \text{E}\left\{ \phantom{ \begin{array}{} \\ \\ \\ \end{array} } \right. \overbrace{ \underbrace{ \begin{array}{|c|c|c|c|c|} \hline ~ & ~ & w_{i1} & ~ & ~ \\ \hline ~ & ~ & w_{i2} & ~ & ~ \\ \hline ~ & ~ & w_{i3} & ~ & ~ \\ \hline \end{array} }_ {\mathbf{W}^\top} }^ {\displaystyle\mathrm{V}} ~~~~~~=~~~~~~ \text{N} \left\{ \phantom{ \begin{array}{} \\ \\ \\ \\ \end{array} } \right. \overbrace{ \underbrace{ \begin{array}{|c|c|c|c|c|} \hline ~ & ~ & \mathbf{u}\mathbf{w}_i & ~ & ~\\ \hline ~ & ~ & ~ & ~ & ~\\ \hline ~ & ~ & ~ & ~ & ~\\ \hline ~ & ~ & ~ & ~ & ~\\ \hline \end{array} }_{\mathrm{softmax}~~ \Rightarrow~~ p_i} }^ {\displaystyle\mathrm{V}} $$

Матриця векторів знаходиться в model[0].weight, а матриця вихідного шару в model[2].weight. Після закінчення навчання, щоб не втратити інформацію, що знаходиться в матриці $\mathbf{W}$, можна усереднити вектори з Embedding і матриці Linear: 

E = model[0].weight                 
W = model[2].weight                 
new_E = 0.5*(E+W)                   # (V, E)

Помилка для Skip-gram Softmax

Виходом моделі повинні бути ймовірності слів. Тому при навчанні використовується nn.CrossEntropyLoss, яка в PyTorch обчислить від отриманих виходів моделі функцію softmax. Так, якщо для $i$-го прикладу необхідно максимізувати ймовірність слова з номером $\hat{y}_i=c$, а виходи моделі дорівнюють $y_{i\alpha}$, то помилка (усереднювана по всіх прикладах батчу) дорівнює:

$$ L(y, c) = -\,w_c\,\log\left( \frac{e^{y_{ic}}}{ \sum_\alpha e^{y_{i\alpha}}}\right). $$ Ваги $w_\alpha$ (за числом слів) можуть підсилювати внесок окремих слів. У нашому випадку будемо підвищувати вагу рідкісніших слів і знижувати вагу часто вживаних слів.

Як ваги можна, наприклад, взяти від'ємний логарифм ймовірності слова $w_\alpha=-\log p(w_\alpha)$ у тексті. Однак, оскільки у нас є багато невеликих документів, скористаємося мірою IDF (inverse document frequency), рівною логарифму оберненої частки документів у яких зустрілося $i$-те слово: $$ \mathrm{idf}(w_i)=1+\log\frac{ N_D}{N_i+1}, $$ де $N_D$ - число документів (у ROCStories - історій), а $N_i$ - число документів у яких зустрілося слово $w_i$: 

    idf = np.zeros( (V_DIM,), dtype=np.float32)    
    for d in docs:                                         # по документах
        for w in set(  w for s in d for w in s.split()  ): # словник документа
            if w in wordID:                                # слово з нашого словника
                idf[wordID[w]["id"]] += 1
                
    idf = ( 1+np.log(len(docs)/(idf+1), dtype=np.float32) )
    weight = idf*(V_DIM/idf.sum())

При створенні помилки моделі необхідно передати їй отримані ваги: 

criterion = nn.CrossEntropyLoss(weight = weight)
У результаті, наприклад, помилка слова "the" буде множитися на 0.12, слова "cat" на 0.62, а "zeus" на 1.27. Вище намальовано графік ваг, як функція номера слова в словнику.

Ваги віднормовані таким чином, щоб їхня сума дорівнювала числу слів у словнику V_DIM. Це дозволяє приблизно порівнювати помилку без ваг (у якій усі ваги дорівнюють 1) і помилку з вагами. Тим не менш слід пам'ятати, що ці помилки (навіть для однієї й тієї самої моделі) будуть відрізнятися. Без ваг помилка зазвичай менша, оскільки моделі достатньо навчитися передбачати тільки високочастотні слова. Наприклад у крапки "." pm = 97442, тобто її ймовірність приблизно дорівнює 0.1. Достатньо завжди передбачати крапку, щоб отримати акуратність 0.1.

Навчання Skip-gram Softmax

Одна епоха навчання зі 100 прикладами в батчі і 45'671'880 загальним числом прикладів на CPU займає близько півтори години. Тому доцільно використовувати для навчання графічний процесор (GPU). 

cpu = torch.device("cpu")  
gpu = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  

model.to(gpu)
weight = weight.to(gpu)

Відповідно в кожному циклі навчання дані батчу також пересилаються в GPU: 
xb = torch.from_numpy( X[it: it+batch_size] ).to(GPU)
yb = torch.from_numpy( Y[it: it+batch_size] ).to(GPU)
       
y = self.model(xb)
Тепер, на одну епоху витрачається близько 20 хвилин.

У випадку ROC Stories початкова помилка, як з вагами, так і без них становить порядка 9.41. Навчена модель з вагами має loss=7.3, а без ваг: 5.3. Така сама ситуація з точністю (помноженою на ваги).

Skip-gram Negative Sampling

Швидший варіант називається skip-gram з negative sampling. Щоб врахувати несиметричність пари $[\mathbf{u},\mathbf{v}]$, будемо для них використовувати різні ембединг-вектори. Знайдемо між ними скалярний добуток і пропустимо його через сигмоїд: $$ p = \frac{1}{1+e^{-\mathbf{u}\cdot\mathbf{v}'}}. $$

Запишемо реалізацію цієї моделі на PyTorch у функціональному вигляді: 
class Skip_gram_Negative1(nn.Module):
    def __init__(self, V_DIM, E_DIM):        
        super(Skip_gram_Negative1, self).__init__()
        
        self.emb1 = nn.Embedding(V_DIM, E_DIM, scale_grad_by_freq=True) 
        self.emb2 = nn.Embedding(V_DIM, E_DIM, scale_grad_by_freq=True) 
        
    def forward(self, x):                            # X з двох колонок
        u  = nn.Flatten() ( self.emb1(x[:,0]) )
        v  = nn.Flatten() ( self.emb2(x[:,1]) )
                        
        dot = u.mul_(v).sum(dim=1)                
        return torch.sigmoid(dot)                    # використовуй BCELoss

Другий спосіб урахування "несиметричності" полягає у використанні єдиного ембедингу, з поворотом другого вектора: $\mathbf{v}'=\mathbf{v}\cdot\mathbf{W}$. Для цього він пропускається через навчуваний лінійний шар Linear з матрицею $\mathbf{W}$ форми (E, E). Потім обчислимо сигмоїд від косинуса між векторами: Таким чином, ймовірність того, що вектори $\mathbf{u}$ і $\mathbf{v}$ виявляться сусідами дорівнює: $$ p = \frac{1}{1+e^{-\mathbf{u}\cdot\mathbf{W}\cdot\mathbf{v}}}. $$ Для векторів $(\mathbf{u},\mathbf{v})$ з одного контексту (паралельних) $p\sim 1$, а для векторів з різних контекстів $p\sim 0$ (обидва вектори беруться з одного Embedding). Навчувана матриця $\mathbf{W}$ відіграє роль "метричного тензора" при обчисленні скалярного добутку. 

class Skip_gram_Negative2(nn.Module):
    def __init__(self, V_DIM, E_DIM):        
        super(Skip_gram_Negative2, self).__init__()
        
        self.emb = nn.Embedding(V_DIM,E_DIM, scale_grad_by_freq=True)
        self.fc  = nn.Linear   (E_DIM, E_DIM)
        
    def forward(self, x):                     
        u  = nn.Flatten() ( self.emb(x[:,0]) )
        v  = nn.Flatten() ( self.emb(x[:,1]) )
        
        dot = (u * self.fc(v)).sum(dim=1)             
        return torch.sigmoid(dot)                    # використовуй BCELoss

Проблема Skip-gram з softmax у тому, що для кожної навчальної пари необхідно обчислювати скалярні добутки $\mathbf{u}\mathbf{w}_i$ з усіма словами в словнику. А це доволі довгий захід. Перемноження матриць (N,E).(E,V) = (N,V) і потім обчислення sum exp (N,V), вимагає порядка N*V*2*E операцій. 

(N,E).(E,V)     -> sum exp (N,V)     

(N,E)*(N,E).sum -> exp (N)           2*N*E

CBOW

У найпростішій версії методу CBOW проводиться усереднення векторів, що входять у модель слів оточення. Результуючий середній вектор пропускається через матрицю лінійного шару Linear з V_DIM нейронами і потім через функцію Softmax (у PyTorch цього робити не треба, якщо використовується помилка CrossEntropyLoss). При цьому ймовірності передбачуваного слова дорівнюють $p_i = e^{\mathbf{u}\mathbf{w}_i}/\sum_j e^{\mathbf{u}\mathbf{v}_j},$ де $\mathbf{u}$ - середній вектор слів оточення.

Праворуч від архітектури мережі намальована форма тензора після виходу з шару Embedding. Модель на PyTorch має вигляд: 

class SBOW(nn.Module):
    def __init__(self, V_DIM, E_DIM):        
        super(SBOW, self).__init__()        
        
        self.emb = nn.Embedding(V_DIM,E_DIM, scale_grad_by_freq=True)
        self.fc  = nn.Linear   (E_DIM, V_DIM)
        
    def forward(self, x):              
        x = self.emb(x)
        x = torch.mean(x, dim=1)
        x = self.fc (x)                
        return x                              # використовуй CrossEntropyLoss

і містить 2,011,005 параметрів (1,000,500 в Embedding і 1,010,505 в Linear). Число слів оточення на число параметрів не впливає.

При підготовці даних, для цільового виходу $Y$ кожного прикладу вказується тільки номер слова (N, ) (sparse-кодування) і використовується помилка torch.nn.CrossEntropyLoss().
Відповідно для вікна з n = 2*WIN входами також складаються з цілих чисел форми (N, n).

Усереднення векторів можна проводити з різними вагами і зробити навчуваним. Для цього замість усереднення вставимо прихований повнозв'язний шар з HIDDEN виходами і relu-функцією активації: 

import torch.nn.functional as F

class SBOW_Hidden(nn.Module):
    def __init__(self, V_DIM, E_DIM, HIDDEN, INPUTS):        
        super(SBOW_Hidden, self).__init__()
        
        self.emb = nn.Embedding(V_DIM, E_DIM, padding_idx=0, scale_grad_by_freq=True)
        self.fc1  = nn.Linear  (INPUTS*E_DIM, HIDDEN)
        self.fc2  = nn.Linear  (HIDDEN, V_DIM)
        
    def forward(self, x):                     
        x = self.emb(x)
        x = nn.Flatten()(x)        
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)                
        return x                                # використовуй CrossEntropyLoss
 
model = SBOW_Hidden(V_DIM, 100, 100, 2*WIN)

Word2Vec і семантичні аналогії

Порівняння методів