7/23 Transformer Blocks

Multi-Head Self-Attention Blocks

In a (Multi-Head Self-) Attention Block, besides the multi-head self-attention operation itself, that we discussed yesterday, we use the tools we learned about on July 20: layer normalization, dropout and residual connections.

1. In the original paper, they used the following formula:

   \[ Y^{(l)}=\mathop{\mathrm{LayerNorm}}\left(X^{(l)} + \mathop{\mathrm{dropout}}\left(\mathop{\mathrm{MHSA}}^{(l)}\left( X^{(l)} \right)\right)\right). \]

2. The above is called Post-LN in [1]. This is because they discovered that what they call Pre-LN:

   \[ Y^{(l)}=X^{(l)} + \mathop{\mathrm{dropout}}\left(\mathop{\mathrm{MHSA}}^{(l)}\left( \mathop{\mathrm{LayerNorm}}\left(X^{(l)}\right) \right)\right). \]

   usually gives better results.

3. Then it was discovered that scale normalization is usually enough:

   \[ Y^{(l)}=X^{(l)} + \mathop{\mathrm{dropout}}\left(\mathop{\mathrm{MHSA}}^{(l)}\left( \mathop{\mathrm{RMSNorm}}\left(X^{(l)}\right) \right)\right). \]

   usually gives better results. Let's call this Pre-RMSN.

Feedforward Blocks

We also need to use a nonlinear transformation along the feature dimension. To this end, we use a so-called feedforward block. Let

\[ \mathop{\mathrm{MLP}}^{(l)}(Y)=a(YW_1^{(l)} + \mathbf b_1^{(l)}) W_2^{(l)} + \mathbf b_2^{(l)} \]

denote an MLP with 1 hidden layer.

Then the Post-LN, Pre-LN, and Pre-RMSN versions of the feedforward block are

\[ X^{(l+1)}=\mathop{\mathrm{LayerNorm}}\left(Y^{(l)} + \mathop{\mathrm{dropout}}\left(\mathop{\mathrm{MLP}}^{(l)}\left( Y^{(l)} \right)\right)\right), \]

\[ X^{(l+1)}=Y^{(l)} + \mathop{\mathrm{dropout}}\left(\mathop{\mathrm{MLP}}^{(l)}\left( \mathop{\mathrm{LayerNorm}}\left(Y^{(l)}\right) \right)\right), \]

and

\[ X^{(l+1)}=Y^{(l)} + \mathop{\mathrm{dropout}}\left(\mathop{\mathrm{MLP}}^{(l)}\left( \mathop{\mathrm{RMSNorm}}\left(Y^{(l)}\right) \right)\right). \]

Transformer Blocks

The composite of a multi-head self attention and a feedforward block is a transformer block. These are the basic building blocks of transformers.

Note that, so far, our architecture is permutation-equivariant. Next time, we'll discuss injecting positional data.

Position Embedding

Note that a transformer layer is permutation equivariant in the sequence dimension. To add token ordering signals, instead of hard-coded rules, we use position embedding vectors.

Absolute Position Embedding

For each sequential position index \(j\), we take a position embedding vector \(\mathbf p_j\in\mathbf R^d\). Then, if in the input text, we have token \(i\) at index \(j\), the embedding vector \(X^{(0)}_j\in\mathbf R^d\), at position \(j\) is the sum \(\mathbf v_i+\mathbf p_j\).

Learned Absolute Position Embeddings

One possible way to define the vectors \(\mathbf p_j\) is to make them trainable vectors, just like the token embeddings \(\mathbf v_i\). We refer to this as learned position embedding. They are used, for example, in BERT [2], the most often used architecture for text and token classification tasks.

Fixed Absolute Position Embeddings

Another possibility is to define the \(\mathbf p_j\) with a formula. We refer to this as fixed position embedding. For example, in the original transformer [3], they use trigonometric functions: the \(k\)-th component of the \(j\)-th position vector is

\[ p_{j, k}=\begin{cases} \sin\left(\frac{j}{10000^{2k / d}}\right) & 2\mid k \\ \cos\left(\frac{j}{10000^{2(k - 1) / d}}\right) & 2 \nmid k \end{cases} \]

They note that this does not have worse performance than learned position embeddings.

Relative Position Embedding

Recall that the sequence entry representations interact through the self-attention operation

\[ \mathop{\mathrm{softmax}}\left(\frac{QK^T}{\sqrt d_k}\right)V \]

where

\[ Q=XW_Q,\,K=XW_K\text{ and }V=XW_V. \]

The idea in relative position embedding is to inject positional information during this quadratic operation.

A general early form [3] is to replace the above formula by

\[ \left( \mathop{\mathrm{softmax}}\left(\frac{QK^T + A_K}{\sqrt d_k}\right) + A_V \right)V \]

where the 3-tensors \(A_K,A_V\in\mathbf R^{s\times s\times d}\) encode positional information.

An important gain for text processing is that this way, we can encode relative distances between tokens, which models language understanding better. Thus, in [4] it is proposed to form the 3-tensors \(A_K\) and \(A_V\) using relative position vectors \(\mathbf p_{i-j}^K,\mathbf p_{i-j}^V\in\mathbf R^d\). They propose to make these learnable vectors.

For efficient implementation, that makes use of the repetitions in the tensors \(A_K\) and \(A_V\), see [4, Subsection 3.3].

References

[1] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang and Tieyan Liu: On Layer Normalization in the Transformer Architecture. Proceedings of the 37th International Conference on Machine Learning, PMLR, vol. 119, pp. 10524--10533, 2020. link

[2] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171–4186, Minneapolis, Minnesota. Association for Computational Linguistics. link

[3] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin: Attention Is All You Need. Advances in Neural Information Processing Systems 30 (NeurIPS 2017) link

[4] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. 2018. Self-Attention with Relative Position Representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers), pages 464–468, New Orleans, Louisiana. Association for Computational Linguistics. link