無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

風雨中的小七發表於2021-02-08
ORM模型

上一章我們聊了聊quick-thought通過幹掉decoder加快訓練, CNN—LSTM用CNN作為Encoder平行計算來提速等方法，這一章看看拋開CNN和RNN，transformer是如何只基於attention對不定長的序列資訊進行提取的。雖然Attention is All you need論文字身是針對NMT翻譯任務的,但transformer作為後續USE/Bert的重要元件，放在embedding裡也沒啥問題。以下基於WMT英翻中的任務實現了transfromer，完整的模型程式碼詳見DSXiangLi-Embedding-transformer

模型元件

讓我們先過一遍Transformer的基礎元件，以文字場景為例，encoder和decoder的輸入是文字序列，每個batch都pad到相同長度然後對每個詞做embedding得到batch * padding_len * emb_size的輸入向量

假設batch=1，Word Embedding維度為512，Encoder的輸入是'Fox hunt rabbit at night', 經過Embedding之後得到1 * 5 * 512的向量，以下的模型元件都服務於如何從這條文字里提取出更多的資訊

Attention

序列資訊提取的一個要點在於如何讓每個詞都考慮到它所在的上下文語境

  • RNN：上下文資訊靠向後/前傳遞，從前往後傳rabbit就能考慮到fox，從後往前傳rabbit就能考慮到night
  • CNN：靠不同kernel_size定義的區域性視窗來獲取context資訊, kernel_size>=3，rabbit就能考慮到所有其他token的資訊
  • Attention：通過計算詞和上下文之間的相關性（廣義），來決定如何把周圍資訊(value)融合(weighted-average)進當前資訊(query），下圖來源Reference5

無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

Transformer在attention的基礎上有兩點改良， 分別是Scaled-dot product attention和multi-head attention。

Scaled-dot product attention

無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

Attention的輸入是三要素query，key和value，通過計算query和Key的相關性，這裡是廣義的相關，可以通過加法/乘法得到權重向量，用權重對value做加權平均作為輸出。‘fox hunt rabbit at night’會計算每個詞對所有詞的相關性，得到[5, 5]的相似度矩陣/權重向量，來對輸入[5, 512]進行加權，得到每個詞在考慮上下文語義後新的向量表達[5, 512]

Transformer在常規的乘法attention的基礎上加入\(d_k\)維度的正則化。這裡\(d_k\)是query和key的特徵維度，在我們的文字場景下是embedding_size[512] 。正則化的原因是避免高維embedding的內積出現超級大的值，導致softmax的gradient非常小。

直觀解釋，假設query和key的每個元素都獨立服從\(\mu=0 \, \sigma^2=1\)的分佈, 那內積\(\sum_{d_k}q_ik_i\)就服從\(\mu=0 \, \sigma^2=d_k\)的分佈,因此需要用\(\sqrt{d_k}\)做正則化，保證內積依舊服從\(\mu=0 \, \sigma=1\)的分佈。

\[Attentino(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}}) \cdot V \]

def scaled_dot_product_attention(key, value, query, mask):
    with tf.variable_scope('scaled_dot_product_attention', reuse=tf.AUTO_REUSE):
        # scalaed weight matrix : batch_size * query_len * key_len
        dk = tf.cast(key.shape.as_list()[-1], tf.float32)# emb_size
        weight = tf.matmul(query, key, transpose_b=True)/(dk**0.5)
        # apply mask: large negative will become 0 in softmax[mask=0 ignore]
        weight += (1-mask) * (-2**32+1)
        # normalize on axis key_len so that score add up to 1
        weight = tf.nn.softmax(weight, axis=-1)
        tf.summary.image("attention", tf.expand_dims(weight[:1], -1))  # add channel dim
        add_layer_summary('attention', weight)
        # weighted value: batch_size * query_len * emb_size
        weighted_value = tf.matmul(weight, value )

        return weighted_value

Mask

上面程式碼中的mask是做什麼的呢？mask決定了Attention對哪些特徵計算權重，transformer的mask有兩種【以下mask=1是保留的部分，0是drop的部分】

其一是padding mask, 讓attention的權重只針對真實文字計算其餘為0。padding mask的dimension是[batch, 1, key_len], 1是預留給query，會在attention中被broadcast成[batch, query_len, key_len]

def seq_mask_gen(input_, params):
    mask = tf.sequence_mask(lengths=tf.to_int32(input_['seq_len']), maxlen=tf.shape(input_['tokens'])[1],
                            dtype=params['dtype'])
    mask = tf.expand_dims(mask, axis=1)
    return mask

如果輸入文字長度分別為3，4，5，都padding到5，padding mask維度是[3,1,5] 如下

無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

其二是future mask只用於decoder，mask每個token後面的序列，保證在預測T+1的token時只會用到T及T以前的資訊，如果不加future mask，預測T+1時就會用到T+1的文字本身，出現feature leakage。

def future_mask_gen(input_, params):
    seq_mask = seq_mask_gen(input_, params) # batch_size * 1 * key_len
    mask = tf.matmul(seq_mask, seq_mask, transpose_a=True) # batch_size * key_len * key_len
    mask = tf.matrix_band_part(mask, num_lower=-1, num_upper=0)
    return mask

還是上面的例子，future mask的維度是[3,5,5] 如下

無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

multi-head attention

無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

這些年對multi-head為啥有效的討論有很多，下面Reference3~8都從不同方面給出了不同的Insight。最開始看multi-head的設計，第一反應是你莫不是在逗我？！你把?掰成8瓣再拼回來告訴我這不是一個?了？？？翻過來倒過去的琢磨感覺這個設計和CNN的filters似乎有些同一個配方，熟悉的味道～也沒啥嚴謹的證明，要是跑偏了也請指正～

假設\(d_{model}=512\), \(head=8\), \(d_k=d_v=d_{model}/head=64\)

multi-head attention的計算過程是，每個head都進行如下操作

  1. 8個head的query，key，value各自過一層線性對映從維度從512->64，權重矩陣\(W \in R^{512*64}\)，這裡和single-head所需的parameter數量一致，single-head是直接做512->512的對映
  2. 對映後的query, key, value做scaled-dot product attention，得到batch * input_len * 64的輸出

\[head_i = Attention(QW_i^Q , K W_i^K , V W_i^V ) \]

再把8個head進行拼接，得到和輸入相同維度\(d_{model}\)的輸出，針對輸出再做一步線性對映就齊活了。

\[MultiHead(Q, K, V ) = Concat(head_1, ..., head_h)\cdot W \]

清楚計算方式我們來看下multi-head和single-head的差異。Single-head每個token, 會用全部512維的embedding來和其他token的embedding計算相關性。這裡其實是存在bottleneck的，因為不論512維的資訊多麼豐富，也只能通過兩兩內積歸一化後得到scaler來表達，多豐富的資訊都會被平均，這一步是存在資訊損失的。既然存在bottleneck，那何不降低\(d_{model}\)的維度，增加head呢。multi-head類似把[5,512]的輸入，先reshape成[5,64,8]，這裡8類似CNN的channel，通過第一步的線性對映我們可能讓每個head（channel）的key，query，value都分別關注不同資訊，可能有的是語法，語義，語序等等，然後用降維到64維的embedding來計算Attention。從而在不增加parameter的條件下提取更多的資訊。

def multi_head_attention(key, value, query, mask, params, mode):
    with tf.variable_scope('multi_head_attention', reuse=tf.AUTO_REUSE):
        d_model = value.shape.as_list()[-1] # emb_size
        # linear projection with dimension unchaangned
        new_key = tf.layers.dense(key, units=d_model, activation=None) # batch_size * key_len * emb_size
        new_value = tf.layers.dense(value, units=d_model, activation=None)
        new_query = tf.layers.dense(query, units=d_model, activation=None)

        # split d_model by num_head and compute attention in parallel
        # (batch_size * num_head) * key_len * (emb_size/num_head)
        new_key = tf.concat(tf.split(new_key, num_or_size_splits=params['num_head'], axis=-1), axis=0)
        new_value = tf.concat(tf.split(new_value, num_or_size_splits=params['num_head'], axis=-1), axis=0)
        new_query = tf.concat(tf.split(new_query, num_or_size_splits=params['num_head'], axis=-1), axis=0)

        # calculate dot-product attention
        weighted_val = scaled_dot_product_attention(new_key, new_value, new_query, tf.tile(mask, [params['num_head'], 1, 1]))

        # concat num_head back
        # (batch_size * num_head) * query_len * (emb_size/num_head) -> batch_size * query_len * emb_size
        weighted_val = tf.concat(tf.split(weighted_val, num_or_size_splits=params['num_head'], axis=0), axis=-1)

        # Linear projection
        weighted_val = tf.layers.dense(weighted_val, units=d_model, activation=None)
        # Do dropout
        weighted_val = tf.layers.dropout(weighted_val, rate=params['dropout_rate'],
                                         training=(mode == tf.estimator.ModeKeys.TRAIN))
        add_layer_summary('raw_multi_head', weighted_val)
        weighted_val = add_and_norm_layer(query, weighted_val)

    return weighted_val

positional encoding

清楚了上面的multi-head的attention會發現有一個問題，就是attention的計算並沒有考慮到詞的相對和絕對位置，這意味著‘fox hunt rabbit at night’和'rabbit hunt fox at night'會有完全一樣的向量表達。Transformer的處理是加入了positional encoding，模型的輸入變為word embedding + positional encoding，因此要求positional encoding和embedding的維度相同都是\(d_{model}\)。

舉個?(以下句子去除停用詞）

  • 句1: Dog sit on chair
  • 句2: Cat is sleeping on floor

要想全面的表達位置資訊，transformer需要滿足以下4個條件

  1. 相對距離: on和chair，on和floor, sit 和on, sleeping和on的相對距離都是1，它們之間的相對距離相同，且和絕對位置以及句子長度無關
  2. 絕對位置：dog和cat都是句子的第一個詞，它們的絕對位置相同, encoding需要一致
  3. 句子長度：encoding需要能夠generalize到訓練樣本中unseen的句子長度

如果我們用[0,句子長度]，步長為1，不滿足條件3，如果測試集出現更長的句子會無法處理。如果用[0,1], 步長為1/句長，不滿足條件2，因為不同長度的句子步長代表的相對距離不一致。讓我們看看Transformer是如何做encoding的

\[PE(pos, i)= \begin{cases} sin(pos /w_k),& \text{if i=2k}\\ cos(pos /w_k),& \text{if i=2k+1} \end{cases} \]

\[w_k = 10000^{2k/d_{model}} \]

因為PE不是trainable變數，所以可以在最開始算好，然後用輸入的position去lookup的，實現如下

def positional_encoding(d_model, max_len, dtype):
    with tf.variable_scope('positional_encoding'):
        encoding_row = np.array([10000**((i-i%2)/d_model) for i in range(d_model)])
        encoding_matrix = np.array([i/encoding_row for i in range(max_len)])

        def sin_cos(row):
            row = [np.cos(val) if i%2 else np.sin(val) for i, val in enumerate(row)]
            return row

        encoding_matrix = np.apply_along_axis(sin_cos, 1, encoding_matrix)
        encoding_matrix = tf.cast(tf.constant(encoding_matrix), dtype)

    return encoding_matrix

還是上面的句子，假設\(d_{model}=4\)，Dog sit on chair的\(PE\)如下

\[\begin{array}{c|c} d_{model} & \text{Dog} & \text{sit} & \text{on} & \text{chair} \\ \hline 0 & sin(\frac{0}{10000^{0/4}}) & sin(\frac{1}{10000^{0/4}}) & sin(\frac{2}{10000^{0/4}}) & sin(\frac{3}{10000^{0/4}}) \\ 1 & cos(\frac{0}{10000^{0/4}}) & cos(\frac{1}{10000^{0/4}}) & cos(\frac{2}{10000^{0/4}}) & cos(\frac{3}{10000^{0/4}}) \\ 2 & sin(\frac{0}{10000^{2/4}}) & sin(\frac{1}{10000^{2/4}}) & sin(\frac{2}{1000^{2/4}}) &sin(\frac{3}{1000^{2/4}}) \\ 3 & cos(\frac{0}{10000^{2/4}}) & cos(\frac{1}{10000^{2/4}}) & cos(\frac{2}{10000^{2/4}}) & cos(\frac{3}{10000^{2/4}}) \\ \end{array} \]

清楚計算方式後，不知道你是不是也有如下的困惑

  • 這sin/cos的設計有何目的？肯定不是為了好看嘛
  • \(w_k\)的計算又是為了啥？encoding不能是個scaler麼？

第一個困惑要看下positional encoding的使用場景。之前提到positional encoding是直接加在word embedding上作為輸出，之後會在計算attention的過程中在兩兩向量內積時被使用，可能會表達類似相對距離越近attention權重越高之類的資訊。因此這裡表達位置和距離是依賴encoding做向量乘法，而使用sin/cos的好處在於位移和絕對位置無關，也就是\(PE(pos+\Delta)=f(\Delta) * PE(pos)\)，詳細推導看這裡Timo Denk's Blog, 以\(d_{model}=2\)為例，線性變換如下

\[\begin{bmatrix} sin(w_k, pos+\Delta) \\ cos(w_k, pos+\Delta)\\ \end{bmatrix}= \begin{bmatrix} cos(w_k,\Delta) & sin(w_k,\Delta) \\ -sin(w_k,\Delta) & cos(w_k,\Delta) \\ \end{bmatrix} \cdot \begin{bmatrix} sin(w_k, pos) \\ cos(w_k, pos)\\ \end{bmatrix} \]

第二個困惑也就是\(w_k\)在embedding維度的計算。有人說是為了和embedding做向量加法，但上面的線性變換隻要有一個[sin,cos]對就能做到，那我把\(R^2\) broadcast到\(R^{d_{model}}\)不成麼， 畢竟PE只是個常量並不trinable。這裡\(w_k\)的計算是隨著k的上升降低了sin/cos的Frequency, PE在不同pos隨i的變化如下圖

無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

看個極端case當\(2k \to d_{model}\)，PE會近似constant，其中\(sin(pos/k) \to 0\), \(cos(pos/k) \to 1\)，和embedding結合來看，部分語義資訊的提取更多依賴位置資訊，自然也存和位置資訊依賴較少或者無關的資訊，在embedding緯度上做差異化的位置資訊表達，可以幫助模型學到這一點～

Add & Norm

Transformer 每個Block之後都會都跟一層Add & Norm,也就是先做residua再做Layer Norm。如果Add & Norm是跟在multi-head Attention之後，這一層的計算便是 Layer_norm(x + multi-head(x))。

def layer_norm(x):
    with tf.variable_scope('layer_normalization', reuse=tf.AUTO_REUSE):
        d_model = x.shape.as_list()[-1]
        epsilon = tf.constant(np.finfo(np.float32).eps)
        mean, variance = tf.nn.moments(x, axes=-1, keep_dims=True)
        x = (x - mean)/((variance + epsilon)**0.5) # do layer norm

        kernel = tf.get_variable('norm_kernel', shape=(d_model,), initializer=tf.ones_initializer())
        bias = tf.get_variable('norm_bias', shape=(d_model,),initializer=tf.zeros_initializer())
        x= tf.multiply(kernel, x) +bias
    return x
    
def add_and_norm_layer(x, sub_layer_x):
    with tf.variable_scope('add_and_norm'):
        x = tf.add(x, sub_layer_x)
        x = layer_norm(x)
    return x

Residual Connection

對於Residual Connection，還是推薦之前在CTR DeepCrossing裡推薦過的一篇文章殘差網路解決了什麼，為什麼有效？。

簡單來說是為了解決網路退化的問題，既隨著網路深度增加，網路的表現先是逐漸增加接近飽和，然後迅速下降。這裡的下降並非指引數增加導致的過擬合，而是理論上如果10層便是最優解，而你的網路有20層，雖然20層包含了10層的資訊，理論上後10層只要做恆等變化把第10層的結果傳遞出去就行，但結果卻變得很差，原因更多懷疑是神經網路較難學習這種恆等變幻。

放在Transformer裡除了以上的作用，也有傳遞底層詞向量和positional encoding資訊的作用，我們既希望通過串聯的Attention來不斷抽取更抽象底層的資訊，但也同時希望向前傳遞Bag of words資訊以及positional encoding攜帶的相對/絕對位置資訊。Anyway殘差結構都帶著更多pratical science的經驗主義，如果有其他的觀點歡迎一起討論喲～

Layer Normalization

LayerNorm也推薦一篇文章詳解深度學習中的Normalization，BN/LN/WN

LayerNorm沒有BatchNorm那麼常用。Batch Norm的假設是所有樣本的同一個特徵（神經元）服從相同的分佈，因此用取樣的樣本（mini-batch）來估計總體在某個特徵上的均值和方差來做歸一化。但BatchNorm對於sequence輸入並不適用，因為不同輸入的序列長度同一個特徵有的樣本有有的沒有，自然不滿足同分布的假設。而LayerNorm的假設是每個樣本的某一層layer是同分布的，因此是每個樣本自身計算stat來做歸一化。

神經元有點抽象，讓我們用傳統ML來舉個?～ 一個預測債券價格的模型有2個特徵：歷史價格和做市商報價。多數情況下的歸一化都是按列進行，所有樣本的歷史價格，做市商報價進行歸一化，對應batch=全樣本的BatchNorm。而LayerNorm對應每個樣本的歷史價格，做市商報價自己進行歸一化，多數情況下因為不同特徵的量綱不同很少做行正則。但如果不同債券的型別不同，多數在100少數在30，而所有特徵都是不同來源的報價，這時對行做正則可能效果更好，因為特徵間分佈比樣本間一致性更高。

回到transformer，這裡的layer Norm是embedding的維度上進行正則化，也就是每個樣本每個token的embedding自身做歸一化。

Feed Forward Layer

每個Multi-head的Attention之後都會跟一個Feed Forward Blok， 是一個兩層的全聯接神經網路, 中間層是relu，既幫助Attention的輸出提取更抽象的資訊，也通過relu過濾無效資訊保留更重要的部分。

\[FFN(h_i) = Relu(h_iw_1 + b_1)w_2 + b_2 \]

def ffn(x, params, mode):
    with tf.variable_scope('ffn', reuse=tf.AUTO_REUSE):
        d_model = x.shape.as_list()[-1]  # emb_size
        y = tf.layers.dense(x, units=params['ffn_hidden'], activation='relu')

        y = tf.layers.dense(y, units=d_model, activation=None)
        y = tf.layers.dropout(y, rate=params['dropout_rate'],
                              training=(mode == tf.estimator.ModeKeys.TRAIN))

        y = add_and_norm_layer(x, y)
    return y

模型實現

愉快的拼樂高時間到，我們來按照以下的模型圖來組合上面的元件，分成encoding和decoding兩個部分。我選了個英文->中文的翻譯任務來實現transformer，完整程式碼詳見DSXiangLi-Embedding-transformer

無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

Encoder

Encoding的輸入是padding的sequence先做詞向量對映得到 batch * pad_len * emb_size的詞向量矩陣, 再加上相同維度的positional encoding向量。計算部分比較簡單是由6個self-attention layer串聯構成。

每個self-attention layer都包括, multi-head attention，encoder source自身既是query也是key和value，過Add&Norm層同時保留變換前和變換後的資訊，再過Feed Forward層做更多的資訊提取，再過Add&Norm。這其中需要注意的便是所有操作的dimension都是\(d_{model}\)，因此輸入緯度不會被改變一直保持到Encoder輸出。

def encode(self, features, mode):
    with tf.variable_scope('encoding', reuse=tf.AUTO_REUSE):
        encoder_input = self.embedding_func(features['tokens'], mode) # batch * seq_len * emb_size
        self_mask = seq_mask_gen(features, self.params)

        for i in range(self.params['encode_attention_layers']):
            with tf.variable_scope('self_attention_layer_{}'.format(i), reuse=tf.AUTO_REUSE):
                encoder_input = multi_head_attention(key=encoder_input, query=encoder_input, value=encoder_input,
                                                     mask=self_mask, params=self.params, mode=mode)

                encoder_input = ffn(encoder_input, self.params, mode)

    return ENCODER_OUTPUT(output=encoder_input, state=encoder_input[:, -1, :])

Decoder

Decoder和encoder一樣也是6個layer串聯。和Encoder相比只是在self-attention和FFN之間多了一層encoder-decoder attention，這時key和value是encoder的輸出，query是decoder在self-attention之後的輸出，學習的是encoder和decoder間的關聯資訊。

Decoding部分略複雜些在於訓練和預測存在差異，原因是訓練會使用teacher forcing用T以前的真實token預測T+1。而預測時真實token未知，因此需要使用loop先預測T=1，拿預測[0,1]去預測T=2，再不斷滾動向前預測

以下是訓練時使用Teacher Forcing的Demo,decoder的輸入文字在source和target要做不同的處理, source第一個token加入\(\lt go\gt\)標記文字開始，如下

無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

這樣預測時預設從\(\lt go\gt\)開始，同時形成錯位用source<=T的token預測T+1的token，剛好對齊target的第T個token，模型預測如下

無所不能的Embedding6 - 跨入Transformer時代～模型詳解&程式碼實現

以下是訓練部分的Decoder

def _decode_helper(self, encoder_output, features, labels, mode):
    decoder_input = self.embedding_func(labels['tokens'], mode)  # batch * seq_len * emb

    self_mask = future_mask_gen(labels, self.params)
    encoder_mask = seq_mask_gen(features, self.params)

    for i in range(self.params['decode_attention_layers']):
        with tf.variable_scope('attention_layer_{}'.format(i), reuse=tf.AUTO_REUSE):
            with tf.variable_scope('self_attention', reuse=tf.AUTO_REUSE):
                decoder_input = multi_head_attention(key=decoder_input, value=decoder_input,
                                                     query=decoder_input, mask=self_mask,
                                                     params=self.params, mode=mode)
            with tf.variable_scope('encode_attention', reuse=tf.AUTO_REUSE):
                decoder_input = multi_head_attention(key=encoder_output.output, value=encoder_output.output,
                                                     query=decoder_input, mask=encoder_mask,
                                                     params=self.params, mode=mode)

            decoder_input = ffn(decoder_input, self.params, mode)

    # use share embedding weight for linear project from emb_size to vocab_size
    logits = tf.matmul(decoder_input, self.embedding, transpose_b=True)  # seq_len * emb_size->seq_len * target_vocab_size

    return DECODER_OUTPUT(output=logits, state=decoder_input, seq_len=labels['seq_len'])

模型訓練

論文中還有不少的訓練細節，例如

  • 每個layer之後都跟了drop out
  • learning rate選取了先上升再下降的noam scheme
  • 樣本按句子長度排序
  • 每個batch保證有近似的單詞數，而非相同的句子數

以及在訓練中發現batch_size太小模型完全不收斂等等。考慮這些細節比較task-specific(有點玄學），感興趣的盆友們可以去看下Martin Popel的Training Tips for the Transformer Model，裡面有更多的細節。後面我們更多隻用到transformer的encoder部分來提取文字資訊，這裡就不多說啦～

Reference

  1. Attention is all you need,
  2. On Layer Normalization in the Transformer Architecture, 2020
  3. Analyzing Multi-Head Self-Attention:
    Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned, 2019
  4. Multi-Head Attention: Collaborate Instead of Concatenate, 2020
  5. What Does BERT Look At? An Analysis of BERT’s Attention, 2019
  6. ON THE RELATIONSHIP BETWEEN SELF-ATTENTION AND CONVOLUTIONAL LAYERS,2020
  7. https://github.com/lena-voita/the-story-of-heads
  8. http://jbcordonnier.com/posts/attention-cnn/
  9. https://towardsdatascience.com/deconstructing-bert-part-2-visualizing-the-inner-workings-of-attention-60a16d86b5c1
  10. http://jalammar.github.io/illustrated-transformer/
  11. 詳解深度學習中的Normalization，BN/LN/WN
  12. Transformer中warm-up和LayerNorm的重要性探究
  13. https://www.zhihu.com/question/341222779/answer/814111138
  14. https://github.com/Kyubyong/transformer

相關文章