Transformer 是一種完全基於注意力機制(attention mechanism)的神經網路架構,最初由 Vaswani 等人於 2017 年提出,後續則有各種不同的改版,如 Conformer (Convolution-augmented Transformer)、Q-former (Querying Transformer)和 Sinkformer (Transformers with Doubly Stochastic Attention)等等,並被廣泛應用於自然語言處理、語音辨識、影像辨識等需要捕捉長距離依賴關係的任務。雖然在應用面上與 LSTM 相似,但不同的是,Transformer 不依賴遞迴結構,而是透過自注意力機制(self-attention)讓序列中的每個位置都能直接與其他所有位置互動,因此在平行運算上具有顯著優勢。本篇章接下來的內容,將基於原始版本進行介紹。

The Transformer is a neural network architecture that is based entirely on the attention mechanism. It was first proposed by Vaswani et al. in 2017, and many versions have followed, such as Conformer (Convolution-augmented Transformer), Q-former (Querying Transformer), and Sinkformer (Transformers with Doubly Stochastic Attention). It is widely used in tasks that need to capture long-distance dependencies, such as natural language processing, speech recognition, and image recognition. It is similar to LSTM in what it is used for, but the difference is that the Transformer does not rely on a recurrent structure. Instead, it uses self-attention, which lets every position in the sequence interact directly with all other positions. So it has a big advantage in parallel computation. The rest of this chapter is based on the original version.

Transformer 的完整架構包含 encoder 與 decoder,我們首先看看 encoder 的部分。

The full Transformer architecture has an encoder and a decoder. Let's look at the encoder first.

Transformer encoder 的核心元件,是由多頭自注意力(Multi-Head Self-Attention, MHSA)與前饋網路(Feed-Forward Network, FFN)所組成。對於給定的輸入序列 X,我們首先分別透過三個線性投影,得到查詢(Query)、鍵(Key)、值(Value)三組矩陣:

The core parts of the Transformer encoder are multi-head self-attention (MHSA) and a feed-forward network (FFN). For a given input sequence X, we first use three linear projections to get three matrices: the Query, the Key, and the Value:

Q = XWQ,K = XWK,V = XWV

接著計算注意力輸出:

Next, we compute the attention output:

Attention(Q, K, V) = softmax(QKT / dk1/2)V

其中 dk 為 K 的維度,除以其平方根是為了避免乘積過大,導致 softmax 梯度消失。多頭注意力則是將上述運算沿特徵維度平均分割成多份來執行,每份使用不同的投影矩陣,最後將結果並排後再做一次線性投影:

Here, dk is the dimension of K. We divide by its square root so the products don't get too large, which would make the softmax gradient vanish. For multi-head attention, we split the operation above into several equal parts along the feature dimension. Each part uses different projection matrices. At the end, we put the results side by side and do one more linear projection:

headi = Attention(XWQi, XWKi, XWVi),

MultiHead(X) = Concat(head1, ..., headh)WO

其中,MultiHead(X) 代表 Q、K、V 皆來自 X。而一個完整的 encoder block,還包含殘差連接(residual connection)與層正規化(Layer Normalization, LN),以及一個逐位置的前饋網路:

Here, MultiHead(X) means that Q, K, and V all come from X. A full encoder block also has a residual connection, layer normalization (LN), and a position-wise feed-forward network:

Z = LN(X + MultiHead(X))

Y = LN(Z + FFN(Z))

其中,FFN 通常是兩層全連接層加上激勵函數:

Here, the FFN is usually two fully connected layers with an activation function:

FFN(Z) = ReLU(ZW1 + b1)W2 + b2

此外,由於 Transformer 沒有遞迴結構,無法從運算過程中得知各位置的順序資訊,因此需要在輸入端加入位置編碼(positional encoding)。原始論文使用正弦與餘弦函數:

Also, the Transformer has no recurrent structure, so it cannot tell the order of the positions from the computation itself. So we need to add positional encoding to the input. The original paper uses sine and cosine functions:

PE[pos, 2i] = sin(pos / 100002i/d)

PE[pos, 2i+1] = cos(pos / 100002i/d)

其中 pos 為時間軸的索引,d 為特徵的總維度,i 為 [0, d/2) 之間的整數。依此填完的 PE 矩陣,其 shape 會與 X 相同,因此我們只要把 X 替換為 X + PE,再餵入 Transformer 即可。

Here, pos is the index along the time axis, d is the total number of feature dimensions, and i is an integer in [0, d/2). The PE matrix filled in this way has the same shape as X. So we only need to replace X with X + PE, and then feed it into the Transformer.

上述的詳細運算步驟,若要用 PyTorch class 的方式表示,則寫法如下:

If we write the detailed steps above as PyTorch classes, it looks like this:

import torch
import torch.nn as nn


class MultiHeadAttention(nn.Module):
	def __init__(self, d_model, num_heads):
		super().__init__()
		self.num_heads = num_heads
		self.d_k = d_model // num_heads
		self.W_q = nn.Linear(d_model, d_model)
		self.W_k = nn.Linear(d_model, d_model)
		self.W_v = nn.Linear(d_model, d_model)
		self.W_o = nn.Linear(d_model, d_model)

	def forward(self, x):
		N, T, d = x.shape
		Q = self.W_q(x).view(N, T, self.num_heads, self.d_k).transpose(1, 2) # (N, h, T, d_k)
		K = self.W_k(x).view(N, T, self.num_heads, self.d_k).transpose(1, 2)
		V = self.W_v(x).view(N, T, self.num_heads, self.d_k).transpose(1, 2)
		scores = Q @ K.transpose(-2, -1) / self.d_k ** 0.5 # (N, h, T, T)
		attn = torch.softmax(scores, dim=-1)
		out = (attn @ V).transpose(1, 2).contiguous().view(N, T, d) # (N, T, d)
		return self.W_o(out)


class PositionalEncoding(nn.Module):
	def __init__(self, d_model):
		super().__init__()
		self.d_model = d_model

	def forward(self, x):
		T = x.size(1)
		pe = torch.zeros(T, self.d_model, device=x.device)
		pos = torch.arange(0, T, device=x.device).unsqueeze(1)
		div = torch.pow(10000.0, torch.arange(0, self.d_model, 2, device=x.device) / self.d_model)
		pe[:, 0::2] = torch.sin(pos / div)
		pe[:, 1::2] = torch.cos(pos / div)
		return x + pe.unsqueeze(0)


class TransformerEncoderBlock(nn.Module):
	def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
		super().__init__()
		self.attn = MultiHeadAttention(d_model, num_heads)
		self.ff = nn.Sequential(
			nn.Linear(d_model, d_ff),
			nn.ReLU(),
			nn.Linear(d_ff, d_model)
		)
		self.norm1 = nn.LayerNorm(d_model)
		self.norm2 = nn.LayerNorm(d_model)
		self.dropout = nn.Dropout(dropout)

	def forward(self, x):
		attn_out = self.attn(x)
		x = self.norm1(x + self.dropout(attn_out))
		ff_out = self.ff(x)
		x = self.norm2(x + self.dropout(ff_out))
		return x


if __name__ == '__main__':
	d_model = 64
	num_heads = 4
	d_ff = 128
	T = 10
	N = 2

	pe = PositionalEncoding(d_model)
	block = TransformerEncoderBlock(d_model, num_heads, d_ff)

	x = torch.randn(N, T, d_model)
	x = pe(x)
	out = block(x)
	print('Input shape:', x.shape)
	print('Output shape:', out.shape)

Transformer 的 decoder block 與 encoder block 結構相似,但多了一個交叉注意力(Cross-Attention)層,其使用的 Q 來自 decoder 自身,K 與 V 則來自 encoder 的輸出。一個完整的 decoder block,依序包含三個子層:

The Transformer decoder block has a structure similar to the encoder block, but it has one more cross-attention layer. In this layer, Q comes from the decoder itself, while K and V come from the output of the encoder. A full decoder block has three sublayers, in this order:

Z1 = LN(Xdec + MultiHead(Xdec))

Z2 = LN(Z1 + CrossAttn(Z1, Xenc))

Y = LN(Z2 + FFN(Z2))

其中 CrossAttn(Z1, Xenc) 表示 Q 來自 Z1,K 與 V 來自 encoder 輸出 Xenc:

Here, CrossAttn(Z1, Xenc) means that Q comes from Z1, while K and V come from the encoder output Xenc:

CrossAttn(Z1, Xenc) = Attention(Z1WQ, XencWK, XencWV)

上述步驟的 PyTorch class 的寫法如下:

The PyTorch classes for the steps above look like this:

class MultiHeadCrossAttention(nn.Module):
	def __init__(self, d_model, num_heads):
		super().__init__()
		self.num_heads = num_heads
		self.d_k = d_model // num_heads
		self.W_q = nn.Linear(d_model, d_model)
		self.W_k = nn.Linear(d_model, d_model)
		self.W_v = nn.Linear(d_model, d_model)
		self.W_o = nn.Linear(d_model, d_model)

	def forward(self, x_dec, x_enc):
		N, T_dec, d = x_dec.shape
		T_enc = x_enc.size(1)
		Q = self.W_q(x_dec).view(N, T_dec, self.num_heads, self.d_k).transpose(1, 2)
		K = self.W_k(x_enc).view(N, T_enc, self.num_heads, self.d_k).transpose(1, 2)
		V = self.W_v(x_enc).view(N, T_enc, self.num_heads, self.d_k).transpose(1, 2)
		scores = Q @ K.transpose(-2, -1) / self.d_k ** 0.5
		attn = torch.softmax(scores, dim=-1)
		out = (attn @ V).transpose(1, 2).contiguous().view(N, T_dec, d)
		return self.W_o(out)


class TransformerDecoderBlock(nn.Module):
	def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
		super().__init__()
		self.self_attn = MultiHeadAttention(d_model, num_heads)
		self.cross_attn = MultiHeadCrossAttention(d_model, num_heads)
		self.ff = nn.Sequential(
			nn.Linear(d_model, d_ff),
			nn.ReLU(),
			nn.Linear(d_ff, d_model)
		)
		self.norm1 = nn.LayerNorm(d_model)
		self.norm2 = nn.LayerNorm(d_model)
		self.norm3 = nn.LayerNorm(d_model)
		self.dropout = nn.Dropout(dropout)

	def forward(self, x_dec, x_enc):
		x = self.norm1(x_dec + self.dropout(self.self_attn(x_dec)))
		x = self.norm2(x + self.dropout(self.cross_attn(x, x_enc)))
		x = self.norm3(x + self.dropout(self.ff(x)))
		return x

上述的 encoder 和 decoder 都各自可以堆疊多層,此時 decoder 的每一層都是從 encoder 最後一層的輸出取 K 和 V,而不是對應層的輸出;也就是說 encoder 跑完所有層之後,只會把最終輸出傳給 decoder 的每一層使用。

Both the encoder and the decoder above can be stacked into many layers. In this case, every decoder layer takes K and V from the output of the last encoder layer, not from the output of the matching layer. In other words, after the encoder runs through all its layers, only the final output is passed to every layer of the decoder.

這個完整的 encoder-decoder 架構,通常會用於生成類的問題,例如將文字從一種語言翻譯成另一種語言。需要注意的是,由於在生成任務的推論過程中,是無法看到未來資訊的;而若要在訓練階段模擬此種行為,則會需要在 decoder 端,利用遮罩(mask)來避免模型看到未來資訊;對於 encoder 端,若你在特定應用中,不希望讓每個位置取用到距離它太遠的資訊,則也會使用 mask 來處理。關於 mask 的細節和使用方式,暫且不列入本篇章的介紹範圍內。

This full encoder-decoder architecture is usually used for generation problems, such as translating text from one language to another. Note that during inference in a generation task, the model cannot see future information. To imitate this during training, we need to use a mask on the decoder side so the model cannot see the future. On the encoder side, if your application does not want each position to use information from positions that are too far away, you can also use a mask. The details and usage of masks are not covered in this chapter for now.

與 LSTM 相仿,PyTorch 也有內建的 class,實務上不需要自己寫。而如果只是要用 Transformer 來進行分類任務,則通常只會取用 encoder 的部分,如下:

Just like LSTM, PyTorch also has a built-in class for this, so in practice you don't need to write it yourself. And if you only want to use a Transformer for a classification task, you usually only use the encoder part, like this:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

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

train_set = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_set = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
print('Data shapes:', train_set.data.shape, test_set.data.shape)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
test_loader = DataLoader(test_set, batch_size=128)


class PositionalEncoding(nn.Module):
	def __init__(self, d_model):
		super().__init__()
		self.d_model = d_model

	def forward(self, x):
		T = x.size(1)
		pe = torch.zeros(T, self.d_model, device=x.device)
		pos = torch.arange(0, T, device=x.device).unsqueeze(1)
		div = torch.pow(10000.0, torch.arange(0, self.d_model, 2, device=x.device) / self.d_model)
		pe[:, 0::2] = torch.sin(pos / div)
		pe[:, 1::2] = torch.cos(pos / div)
		return x + pe.unsqueeze(0)


class TransformerClassifier(nn.Module):
	def __init__(self, d_model=28, num_heads=2, d_ff=64, dropout=0.1):
		super().__init__()
		self.pe = PositionalEncoding(d_model)
		self.encoder = nn.TransformerEncoderLayer(
			d_model=d_model,
			nhead=num_heads,
			dim_feedforward=d_ff,
			dropout=dropout,
			batch_first=True
		)
		self.linear = nn.Linear(d_model, 10)

	def forward(self, x):
		x = x.squeeze(1)
		x = self.pe(x)
		x = self.encoder(x)
		return self.linear(x[:, -1, :])


model = TransformerClassifier().to(device)
optim = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

model.train()
for i in range(10):
	print(f'Epoch {i+1}/10')
	for X_batch, y_batch in train_loader:
		loss = criterion(model(X_batch.to(device)), y_batch.to(device))
		optim.zero_grad()
		loss.backward()
		optim.step()

model.eval()
correct = sum(
	(model(X.to(device)).argmax(dim=1) == y.to(device)).sum().item()
	for X, y in test_loader
)
print('Accuracy: {:.2f}%'.format(100 * correct / len(test_set)))

在上述範例中,我們仿照 LSTM 篇章的做法,只取最後一個時間步的輸出來進行後續的分類,你當然還是可以嘗試其他做法。此外,範例中只使用了一層 encoder,而堆疊多層的做法,就留給同學做練習。

In the example above, we follow what we did in the LSTM chapter. We only take the output of the last time step for the later classification. Of course, you can still try other ways. Also, the example only uses one encoder layer. We leave stacking many layers as an exercise for you.

Transformer 還有其他不少的著名使用案例,如:

The Transformer has quite a few other famous use cases, such as:

需要注意的是,文字的輸入和輸出,不同於本篇章或前面各篇章,對影像或數值的處理。雖然也有少數做法會直接使用每個字元在電腦中的編碼,轉成 one-hot vector 作為輸入或者輸出;但比較主流的做法,是把文字切成小段並依此建立字典,才轉成 one-hot vector 甚至是 embedding 來使用。這個切段的動作稱為 tokenization,且不同的預訓練模型通常會搭配不同的字典,使用時必須互相對應。有興趣了解細節的同學,可以自行尋找相關資料研讀。

Note that text input and output are different from how we handle images or numbers in this chapter and the earlier chapters. A few methods do directly use the computer encoding of each character, and turn it into a one-hot vector as the input or output. But the more common way is to cut the text into small pieces and build a dictionary from them, and only then turn them into one-hot vectors, or even embeddings. This cutting step is called tokenization. Different pre-trained models usually come with different dictionaries, and you must use the matching one. If you are interested in the details, you can look up related materials and study them by yourself.

前述的「預訓練模型」,是指他人用大量資料訓練好的模型,以供你進行個人化的調整;相當於你可以站在巨人的肩膀上來處理特定的問題,不用從頭開始爬山。而調整模型時,若認為整個權重矩陣過於龐大不易訓練,你也可以凍結原始權重,只對其 low-rank 版本訓練修正量,這種技巧稱為 low-rank Adaptation (LoRA)。預訓練模型不限於只在 Transformer 上適用,例如影像處理的 VGG 和 ResNet 系列等模型,都有預訓練版本的釋出;而實務上,Transformer 以外的模型,則基於效益等方面的考量,因此不常使用 LoRA。

The "pre-trained model" mentioned above means a model that other people have trained with a large amount of data, so you can adjust it for your own use. This is like standing on the shoulders of giants to solve a specific problem, without having to climb the hill from the bottom. When adjusting a model, if you think the whole weight matrix is too big to train easily, you can also freeze the original weights and only train a correction in a low-rank form. This technique is called low-rank Adaptation (LoRA). Pre-trained models are not only for Transformers. For example, image models such as the VGG and ResNet series also have pre-trained versions released. In practice, for models other than Transformers, LoRA is not used often, because of cost-benefit and other reasons.