10.3. Kiến trúc Transformer

Trong các chương trước, ta đã đề cập đến các kiến trúc mạng nơ-ron quan trọng như mạng nơ-ron tích chập (CNN) và mạng nơ-ron hồi tiếp (RNN). Ưu nhược điểm của hai kiến trúc mạng này có thể được tóm tắt như sau:

  • Các mạng CNN có thể dễ dàng được thực hiện song song ở một tầng nhưng không có khả năng nắm bắt các phụ thuộc chuỗi có độ dài biến thiên.
  • Các mạng RNN có khả năng nắm bắt các thông tin cách xa nhau trong chuỗi có độ dài biến thiên, nhưng không thể thực hiện song song trong một chuỗi.

Để kết hợp các ưu điểm của CNN và RNN, [Vaswani et al., 2017] đã thiết kế một kiến trúc mới bằng cách sử dụng cơ chế tập trung. Kiến trúc này gọi là Transformer, song song hóa bằng cách học chuỗi hồi tiếp với cơ chế tập trung, đồng thời mã hóa vị trí của từng phần tử trong chuỗi. Kết quả là ta có một mô hình tương thích với thời gian huấn luyện ngắn hơn đáng kể.

Tương tự như mô hình seq2seq trong Section 9.7, Transformer cũng dựa trên kiến trúc mã hóa-giải mã. Tuy nhiên, nó thay thế các tầng hồi tiếp trong seq2seq bằng các tầng tập trung đa đầu (multi-head attention), kết hợp thông tin vị trí thông qua biểu diễn vị trí (positional encoding) và áp dụng chuẩn hóa tầng (layer normalization). Fig. 10.3.1 sẽ so sánh cấu trúc của Transformer và seq2seq.

Nhìn chung, hai mô hình này khá giống nhau: các embedding của chuỗi nguồn được đưa vào \(n\) khối lặp lại. Đầu ra của khối mã hóa cuối cùng sau đó được sử dụng làm bộ nhớ tập trung cho bộ giải mã. Tương tự, các embedding của chuỗi đích được đưa vào \(n\) khối lặp lại trong bộ giải mã. Ta thu được đầu ra cuối cùng bằng cách áp dụng một tầng dày đặc có kích thước bằng kích thước bộ từ vựng lên các đầu ra của khối giải mã cuối cùng.

../_images/transformer.svg

Fig. 10.3.1 Kiến trúc Transformer.

Mặt khác, Transformer khác với mô hình seq2seq sử dụng cơ chế tập trung như sau:

  1. Khối Transformer: một tầng hồi tiếp trong seq2seq được thay bằng một Khối Transformer. Với bộ mã hóa, khối này chứa một tầng tập trung đa đầu và một mạng truyền xuôi theo vị trí (position-wise feed-forward network) gồm hai tầng dày đặc. Đối với bộ giải mã, khối này có thêm một tầng tập trung đa đầu khác để nhận vào trạng thái bộ mã hóa.
  2. Cộng và chuẩn hóa: đầu vào và đầu ra của cả tầng tập trung đa đầu hoặc mạng truyền xuôi theo vị trí được xử lý bởi hai tầng “cộng và chuẩn hóa” bao gồm cấu trúc phần dư và tầng chuẩn hóa theo tầng (layer normalization).
  3. Biễu diễn vị trí: do tầng tự tập trung không phân biệt thứ tự phần tử trong một chuỗi, nên tầng biễu diễn vị trí được sử dụng để thêm thông tin vị trí vào từng phần tử trong chuỗi.

Tiếp theo, chúng ta sẽ tìm hiểu từng thành phần trong Transformer để có thể xây dựng một mô hình dịch máy.

from d2l import mxnet as d2l
import math
from mxnet import autograd, np, npx
from mxnet.gluon import nn
npx.set_np()

10.3.1. Tập trung Đa đầu

Trước khi thảo luận về tầng tập trung đa đầu, hãy cùng tìm hiểu qua về kiến trúc tự tập trung. Giống như các mô hình tập trung bình thường, mô hình tự tập trung cũng có câu truy vấn, khóa và giá trị nhưng chúng được sao chép từ các phần tử trong chuỗi đầu vào. Như minh họa trong Fig. 10.3.2, tầng tự tập trung trả về một đầu ra tuần tự có cùng độ dài với đầu vào. So với tầng hồi tiếp, các phần tử đầu ra của tầng tự tập trung có thể được tính toán song song, do đó việc xây dựng các đoạn mã tốc độ cao khá dễ dàng.

../_images/self-attention.svg

Fig. 10.3.2 Kiến trúc tự tập trung.

Tầng tập trung đa đầu bao gồm \(h\) đầu là các tầng tự tập trung song song. Trước khi đưa vào mỗi đầu, ta chiếu các câu truy vấn, khóa và giá trị qua ba tầng dày đặc với kích thước ẩn lần lượt là \(p_q\), \(p_k\) và \(p_v\). Đầu ra của \(h\) đầu này được nối lại và sau đó được xử lý bởi một tầng dày đặc cuối cùng.

../_images/multi-head-attention.svg

Fig. 10.3.3 Tập trung đa đầu

Giả sử chiều của câu truy vấn, khóa và giá trị lần lượt là \(d_q\), \(d_k\)\(d_v\). Khi đó, tại mỗi đầu \(i=1,\ldots, h\), ta có thể học các tham số \(\mathbf W_q^{(i)}\in\mathbb R^{p_q\times d_q}\), \(\mathbf W_k^{(i)}\in\mathbb R^{p_k\times d_k}\), và \(\mathbf W_v^{(i)}\in\mathbb R^{p_v\times d_v}\). Do đó, đầu ra tại mỗi đầu là

(10.3.1)\[\mathbf o^{(i)} = \textrm{attention}(\mathbf W_q^{(i)}\mathbf q, \mathbf W_k^{(i)}\mathbf k,\mathbf W_v^{(i)}\mathbf v),\]

trong đó \(\textrm{attention}\) có thể là bất kỳ tầng tập trung nào, chẳng hạn như DotProductAttentionMLPAttention trong Section 10.1.

Sau đó, \(h\) đầu ra với độ dài \(p_v\) tại mỗi đầu được nối với nhau thành đầu ra có độ dài \(h p_v\), rồi được đưa vào tầng dày đặc cuối cùng với \(d_o\) nút ẩn. Các trọng số của tầng dày đặc này được ký hiệu là \(\mathbf W_o\in\mathbb R^{d_o\times h p_v}\). Do đó, đầu ra cuối cùng của tầng tập trung đa đầu sẽ là

(10.3.2)\[\begin{split}\mathbf o = \mathbf W_o \begin{bmatrix}\mathbf o^{(1)}\\\vdots\\\mathbf o^{(h)}\end{bmatrix}.\end{split}\]

Bây giờ chúng ta có thể lập trình tầng tập trung đa đầu. Giả sử tầng tập trung đa đầu có số đầu là num_heads \(=h\) và các tầng dày đặc cho câu truy vấn, khóa và giá trị có kích thước ẩn giống nhau num_hiddens \(=p_q=p_k=p_v\). Ngoài ra, do tầng tập trung đa đầu giữ nguyên kích thước chiều đầu vào, kích thước đặc trưng đầu ra cũng là \(d_o =\) num_hiddens.

# Saved in the d2l package for later use
class MultiHeadAttention(nn.Block):
    def __init__(self, num_hiddens, num_heads, dropout, use_bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
        self.W_k = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
        self.W_v = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
        self.W_o = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)

    def forward(self, query, key, value, valid_len):
        # For self-attention, query, key, and value shape:
        # (batch_size, seq_len, dim), where seq_len is the length of input
        # sequence. valid_len shape is either (batch_size, ) or
        # (batch_size, seq_len).

        # Project and transpose query, key, and value from
        # (batch_size, seq_len, num_hiddens) to
        # (batch_size * num_heads, seq_len, num_hiddens / num_heads).
        query = transpose_qkv(self.W_q(query), self.num_heads)
        key = transpose_qkv(self.W_k(key), self.num_heads)
        value = transpose_qkv(self.W_v(value), self.num_heads)

        if valid_len is not None:
            # Copy valid_len by num_heads times
            if valid_len.ndim == 1:
                valid_len = np.tile(valid_len, self.num_heads)
            else:
                valid_len = np.tile(valid_len, (self.num_heads, 1))

        # For self-attention, output shape:
        # (batch_size * num_heads, seq_len, num_hiddens / num_heads)
        output = self.attention(query, key, value, valid_len)

        # output_concat shape: (batch_size, seq_len, num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

Dưới đây là định nghĩa của hai hàm chuyển vị transpose_qkvtranspose_output. Hai hàm này là nghịch đảo của nhau.

# Saved in the d2l package for later use
def transpose_qkv(X, num_heads):
    # Input X shape: (batch_size, seq_len, num_hiddens).
    # Output X shape:
    # (batch_size, seq_len, num_heads, num_hiddens / num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # X shape: (batch_size, num_heads, seq_len, num_hiddens / num_heads)
    X = X.transpose(0, 2, 1, 3)

    # output shape: (batch_size * num_heads, seq_len, num_hiddens / num_heads)
    output = X.reshape(-1, X.shape[2], X.shape[3])
    return output


# Saved in the d2l package for later use
def transpose_output(X, num_heads):
    # A reversed version of transpose_qkv
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.transpose(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

Hãy cùng kiểm tra mô hình MultiHeadAttention qua một ví dụ đơn giản. Tạo một tầng tập trung đa đầu với kích thước ẩn \(d_o = 100\), đầu ra sẽ có cùng kích thước batch và độ dài chuỗi với đầu vào, nhưng có kích thước chiều cuối cùng bằng num_hiddens \(= 100\).

cell = MultiHeadAttention(90, 9, 0.5)
cell.initialize()
X = np.ones((2, 4, 5))
valid_len = np.array([2, 3])
cell(X, X, X, valid_len).shape
(2, 4, 90)

10.3.2. Mạng truyền Xuôi theo Vị trí

Một thành phần quan trọng khác trong Khối Transformer là mạng truyền xuôi theo vị trí (position-wise feed-forward network). Nó chấp nhận đầu vào \(3\) chiều với kích thước là (kích thước batch, độ dài chuỗi, kích thước đặc trưng). Mạng truyền xuôi theo vị trí bao gồm hai tầng dày đặc áp dụng trên chiều cuối cùng của đầu vào. Vì hai tầng dày đặc này cùng được sử dụng cho từng vị trí trong chuỗi, nên ta gọi là mạng truyền xuôi theo vị trí. Cách làm này tương đương với việc áp dụng hai tầng tích chập \(1 \times 1\).

Lớp PositionWiseFFN dưới đây lập trình mạng truyền xuôi theo vị trí với hai tầng dày đặc có kích thước ẩn lần lượt là ffn_num_hiddenspw_num_outputs.

# Saved in the d2l package for later use
class PositionWiseFFN(nn.Block):
    def __init__(self, ffn_num_hiddens, pw_num_outputs, **kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        self.dense1 = nn.Dense(ffn_num_hiddens, flatten=False,
                               activation='relu')
        self.dense2 = nn.Dense(pw_num_outputs, flatten=False)

    def forward(self, X):
        return self.dense2(self.dense1(X))

Tương tự như tầng tập trung đa đầu, mạng truyền xuôi theo vị trí sẽ chỉ thay đổi kích thước chiều cuối cùng của đầu vào — tức kích thước của đặc trưng. Ngoài ra, nếu hai phần tử trong chuỗi đầu vào giống hệt nhau, thì hai đầu ra tương ứng cũng sẽ giống hệt nhau.

ffn = PositionWiseFFN(4, 8)
ffn.initialize()
ffn(np.ones((2, 3, 4)))[0]
array([[ 9.15348239e-04, -7.27669394e-04,  1.14063594e-04,
        -8.76279722e-04, -1.02867256e-03,  8.02748313e-04,
        -4.53725770e-05,  2.15598906e-04],
       [ 9.15348239e-04, -7.27669394e-04,  1.14063594e-04,
        -8.76279722e-04, -1.02867256e-03,  8.02748313e-04,
        -4.53725770e-05,  2.15598906e-04],
       [ 9.15348239e-04, -7.27669394e-04,  1.14063594e-04,
        -8.76279722e-04, -1.02867256e-03,  8.02748313e-04,
        -4.53725770e-05,  2.15598906e-04]])

10.3.3. Cộng và Chuẩn hóa

Trong kiến trúc Transformer, tầng “cộng và chuẩn hóa” cũng đóng vai trò thiết yếu trong việc kết nối đầu vào và đầu ra của các tầng khác một cách trơn tru. Cụ thể, ta thêm một cấu trúc phần dư và tầng chuẩn hóa theo tầng sau tầng tập trung đa đầu và mạng truyền xuôi theo vị trí. Chuẩn hóa theo tầng khá giống với chuẩn hóa theo batch trong Section 7.5. Một điểm khác biệt là giá trị trung bình và phương sai của tầng chuẩn hóa này được tính theo chiều cuối cùng, tức X.mean(axis=-1), thay vì theo chiều đầu tiên (theo batch) X.mean(axis=0) . Chuẩn hóa tầng ngăn không cho phạm vi giá trị trong các tầng thay đổi quá nhiều, giúp huấn luyện nhanh hơn và khái quát hóa tốt hơn.

MXNet có cả LayerNormBatchNorm được lập trình trong khối nn. Hãy cùng xem sự khác biệt giữa chúng qua ví dụ dưới đây.

layer = nn.LayerNorm()
layer.initialize()
batch = nn.BatchNorm()
batch.initialize()
X = np.array([[1, 2], [2, 3]])
# Compute mean and variance from X in the training mode
with autograd.record():
    print('layer norm:', layer(X), '\nbatch norm:', batch(X))
layer norm: [[-0.99998  0.99998]
 [-0.99998  0.99998]]
batch norm: [[-0.99998 -0.99998]
 [ 0.99998  0.99998]]

Bây giờ hãy cùng lập trình khối AddNorm. AddNorm nhận hai đầu vào \(X\)\(Y\). Ta có thể coi \(X\) là đầu vào ban đầu trong mạng phần dư và \(Y\) là đầu ra từ tầng tập trung đa đầu hoặc mạng truyền xuôi theo vị trí. Ngoài ra, ta cũng sẽ áp dụng dropout lên \(Y\) để điều chuẩn.

# Saved in the d2l package for later use
class AddNorm(nn.Block):
    def __init__(self, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)
        self.ln = nn.LayerNorm()

    def forward(self, X, Y):
        return self.ln(self.dropout(Y) + X)

Do có kết nối phần dư, \(X\)\(Y\) sẽ có kích thước giống nhau.

add_norm = AddNorm(0.5)
add_norm.initialize()
add_norm(np.ones((2, 3, 4)), np.ones((2, 3, 4))).shape
(2, 3, 4)

10.3.4. Biểu diễn Vị trí

Không giống như tầng hồi tiếp, cả tầng tập trung đa đầu và mạng truyền xuôi theo vị trí đều tính toán đầu ra cho từng phần tử trong chuỗi một cách độc lập. Điều này cho phép song song hóa công việc tính toán nhưng lại không mô hình hóa được thông tin tuần tự trong chuỗi đầu vào. Để nắm bắt các thông tin tuần tự một cách hiệu quả, mô hình Transformer sử dụng biểu diễn vị trí (positional encoding) để duy trì thông tin vị trí của chuỗi đầu vào.

Cụ thể, giả sử \(X\in\mathbb R^{l\times d}\) là embedding của mẫu đầu vào, trong đó \(l\) là độ dài chuỗi và \(d\) là kích thước embedding. Tầng biểu diễn vị trí sẽ mã hóa vị trí \(P\in\mathbb R^{l\times d}\) của X và trả về đầu ra \(P+X\).

Vị trí \(P\) là ma trận 2 chiều, trong đó \(i\) là thứ tự trong câu, \(j\) là vị trí theo chiều embedding. Bằng cách này, mỗi vị trí trong chuỗi ban đầu được biểu biễn bởi hai phương trình dưới đây:

(10.3.3)\[P_{i, 2j} = \sin(i/10000^{2j/d}),\]
(10.3.4)\[\quad P_{i, 2j+1} = \cos(i/10000^{2j/d}),\]

với \(i=0,\ldots, l-1\)\(j=0,\ldots,\lfloor(d-1)/2\rfloor\).

Fig. 10.3.4 minh họa biểu diễn vị trí.

../_images/positional_encoding.svg

Fig. 10.3.4 Biểu diễn vị trí.

# Saved in the d2l package for later use
class PositionalEncoding(nn.Block):
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        # Create a long enough P
        self.P = np.zeros((1, max_len, num_hiddens))
        X = np.arange(0, max_len).reshape(-1, 1) / np.power(
            10000, np.arange(0, num_hiddens, 2) / num_hiddens)
        self.P[:, :, 0::2] = np.sin(X)
        self.P[:, :, 1::2] = np.cos(X)

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].as_in_ctx(X.ctx)
        return self.dropout(X)

Bây giờ chúng ta kiểm tra lớp PositionalEncoding ở trên bằng một mô hình đơn giản cho 4 chiều. Có thể thấy, chiều thứ 4 và chiều thứ 5 có cùng tần số nhưng khác độ dời, còn chiều thứ 6 và 7 có tần số thấp hơn.

pe = PositionalEncoding(20, 0)
pe.initialize()
Y = pe(np.zeros((1, 100, 20)))
d2l.plot(np.arange(100), Y[0, :, 4:8].T, figsize=(6, 2.5),
         legend=["dim %d" % p for p in [4, 5, 6, 7]])
../_images/output_transformer_vn_e7c020_21_0.svg

10.3.5. Bộ Mã hóa

Với các thành phần thiết yếu trên, hãy xây dựng bộ mã hóa cho Transformer. Bộ mã hóa này chứa một tầng tập trung đa đầu, một mạng truyền xuôi theo vị trí và hai khối kết nối “cộng và chuẩn hóa”. Trong mã nguồn, có thể thấy cả tầng tập trung và mạng truyền xuôi theo vị trí trong EncoderBlock đều có đầu ra với kích thước là num_hiddens. Điều này là do kết nối phần dư trong quá trình “cộng và chuẩn hóa”, khi ta cần cộng đầu ra của hai khối này với giá trị đầu vào của chúng.

# Saved in the d2l package for later use
class EncoderBlock(nn.Block):
    def __init__(self, num_hiddens, ffn_num_hiddens, num_heads, dropout,
                 use_bias=False, **kwargs):
        super(EncoderBlock, self).__init__(**kwargs)
        self.attention = MultiHeadAttention(num_hiddens, num_heads, dropout,
                                            use_bias)
        self.addnorm1 = AddNorm(dropout)
        self.ffn = PositionWiseFFN(ffn_num_hiddens, num_hiddens)
        self.addnorm2 = AddNorm(dropout)

    def forward(self, X, valid_len):
        Y = self.addnorm1(X, self.attention(X, X, X, valid_len))
        return self.addnorm2(Y, self.ffn(Y))

Nhờ kết nối phần dư, khối mã hóa sẽ không thay đổi kích thước đầu vào. Nói cách khác, giá trị num_hiddens phải bằng kích thước chiều cuối cùng của đầu vào. Trong ví dụ đơn giản dưới đây, num_hiddens \(= 24\),ffn_num_hiddens \(=48\), num_heads \(= 8\)dropout \(= 0.5\).

X = np.ones((2, 100, 24))
encoder_blk = EncoderBlock(24, 48, 8, 0.5)
encoder_blk.initialize()
encoder_blk(X, valid_len).shape
(2, 100, 24)

Bây giờ hãy lập trình bộ mã hóa của Transformer hoàn chỉnh. Trong bộ mã hóa, \(n\) khối EncoderBlock được xếp chồng lên nhau. Nhờ có kết nối phần dư, tầng embedding và đầu ra khối Transformer đều có kích thước là \(d\). Cũng lưu ý rằng ta nhân các embedding với \(\sqrt{d}\) để tránh trường hợp giá trị này quá nhỏ.

# Saved in the d2l package for later use
class TransformerEncoder(d2l.Encoder):
    def __init__(self, vocab_size, num_hiddens, ffn_num_hiddens,
                 num_heads, num_layers, dropout, use_bias=False, **kwargs):
        super(TransformerEncoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for _ in range(num_layers):
            self.blks.add(
                EncoderBlock(num_hiddens, ffn_num_hiddens, num_heads, dropout, use_bias))

    def forward(self, X, valid_len, *args):
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        for blk in self.blks:
            X = blk(X, valid_len)
        return X

Hãy tạo một bộ mã hóa với hai khối mã hóa Transformer với siêu tham số như ví dụ trên. Ngoài ra, ta đặt hai tham số vocab_size bằng \(200\)num_layers bằng \(2\).

encoder = TransformerEncoder(200, 24, 48, 8, 2, 0.5)
encoder.initialize()
encoder(np.ones((2, 100)), valid_len).shape
(2, 100, 24)

10.3.6. Bộ Giải mã

Khối giải mã của Transformer gần tương tự như khối mã hóa. Tuy nhiên, bên cạnh hai tầng con (tập trung đa đầu và biểu diễn vị trí), khối giải mã còn chứa tầng tập trung đa đầu áp dụng lên đầu ra của bộ mã hóa. Các tầng con này cũng được kết nối bằng các tầng “cộng và chuẩn hóa”, gồm kết nối phần dư và chuẩn hóa theo tầng.

Cụ thể, tại bước thời gian \(t\), giả sử đầu vào hiện tại là câu truy vấn \(\mathbf x_t\). Như minh họa trong Fig. 10.3.5, các khóa và giá trị của tầng tập trung gồm có câu truy vấn ở bước thời gian hiện tại và tất cả các câu truy vấn ở các bước thời gian trước \(\mathbf x_1, \ldots, \mathbf x_{t-1}\).

../_images/self-attention-predict.svg

Fig. 10.3.5 Dự đoán ở bước thời gian \(t\) của một tầng tự tập trung.

Trong quá trình huấn luyện, đầu ra của câu truy vấn \(t\) có thể quan sát được tất cả các cặp khóa - giá trị trước đó. Điều này dẫn đến hai cách hoạt động khác nhau của mạng khi huấn luyện và dự đoán. Vì thế, trong lúc dự đoán chúng ta có thể loại bỏ những thông tin không cần thiết bằng cách chỉ định độ dài hợp lệ là \(t\) cho câu truy vấn thứ \(t\).

class DecoderBlock(nn.Block):
    # i means it is the i-th block in the decoder
    def __init__(self, num_hiddens, ffn_num_hiddens, num_heads,
                 dropout, i, **kwargs):
        super(DecoderBlock, self).__init__(**kwargs)
        self.i = i
        self.attention1 = MultiHeadAttention(num_hiddens, num_heads, dropout)
        self.addnorm1 = AddNorm(dropout)
        self.attention2 = MultiHeadAttention(num_hiddens, num_heads, dropout)
        self.addnorm2 = AddNorm(dropout)
        self.ffn = PositionWiseFFN(ffn_num_hiddens, num_hiddens)
        self.addnorm3 = AddNorm(dropout)

    def forward(self, X, state):
        enc_outputs, enc_valid_len = state[0], state[1]
        # state[2][i] contains the past queries for this block
        if state[2][self.i] is None:
            key_values = X
        else:
            key_values = np.concatenate((state[2][self.i], X), axis=1)
        state[2][self.i] = key_values
        if autograd.is_training():
            batch_size, seq_len, _ = X.shape
            # Shape: (batch_size, seq_len), the values in the j-th column
            # are j+1
            valid_len = np.tile(np.arange(1, seq_len+1, ctx=X.ctx),
                                   (batch_size, 1))
        else:
            valid_len = None

        X2 = self.attention1(X, key_values, key_values, valid_len)
        Y = self.addnorm1(X, X2)
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_len)
        Z = self.addnorm2(Y, Y2)
        return self.addnorm3(Z, self.ffn(Z)), state

Tương tự như khối mã hóa của Transformer, num_hiddens của khối giải mã phải bằng với kích thước chiều cuối cùng của \(X\).

decoder_blk = DecoderBlock(24, 48, 8, 0.5, 0)
decoder_blk.initialize()
X = np.ones((2, 100, 24))
state = [encoder_blk(X, valid_len), valid_len, [None]]
decoder_blk(X, state)[0].shape
(2, 100, 24)

Bộ giải mã hoàn chỉnh của Transformer được lập trình tương tự như bộ mã hóa, ngoại trừ chi tiết tầng kết nối đầy đủ được thêm vào để tính điểm tin cậy của đầu ra.

Hãy lập trình bộ giải mã đầy đủ TransformerDecoder. Ngoài các siêu tham số thường gặp như vocab_sizenum_hiddens, bộ giải mã cần thêm kích thước đầu ra của bộ mã hóa enc_outputs và độ dài hợp lệ env_valid_len.

class TransformerDecoder(d2l.Decoder):
    def __init__(self, vocab_size, num_hiddens, ffn_num_hiddens,
                 num_heads, num_layers, dropout, **kwargs):
        super(TransformerDecoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add(
                DecoderBlock(num_hiddens, ffn_num_hiddens, num_heads,
                             dropout, i))
        self.dense = nn.Dense(vocab_size, flatten=False)

    def init_state(self, enc_outputs, env_valid_len, *args):
        return [enc_outputs, env_valid_len, [None]*self.num_layers]

    def forward(self, X, state):
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        for blk in self.blks:
            X, state = blk(X, state)
        return self.dense(X), state

10.3.7. Huấn luyện

Cuối cùng, chúng ta có thể xây dựng một mô hình mã hóa - giải mã với kiến ​​trúc Transformer. Tương tự như mô hình seq2seq áp dụng cơ chế tập trung trong Section 10.2, chúng ta sử dụng các siêu tham số sau: hai khối Transformer có kích thước embedding và kích thước đầu ra đều là \(32\). Bên cạnh đó, chúng ta sử dụng \(4\) đầu trong tầng tập trung và đặt kích thước ẩn bằng hai lần kích thước đầu ra.

num_hiddens, num_layers, dropout, batch_size, num_steps = 32, 2, 0.0, 64, 10
lr, num_epochs, ctx = 0.005, 100, d2l.try_gpu()
ffn_num_hiddens, num_heads = 64, 4

src_vocab, tgt_vocab, train_iter = d2l.load_data_nmt(batch_size, num_steps)

encoder = TransformerEncoder(
    len(src_vocab), num_hiddens, ffn_num_hiddens, num_heads, num_layers,
    dropout)
decoder = TransformerDecoder(
    len(src_vocab), num_hiddens, ffn_num_hiddens, num_heads, num_layers,
    dropout)
model = d2l.EncoderDecoder(encoder, decoder)
d2l.train_s2s_ch9(model, train_iter, lr, num_epochs, ctx)
loss 0.033, 3396.2 tokens/sec on gpu(0)
../_images/output_transformer_vn_e7c020_37_1.svg

Dựa trên thời gian huấn luyện và độ chính xác, ta thấy Transformer chạy nhanh hơn trên mỗi epoch và hội tụ nhanh hơn ở giai đoạn đầu so với seq2seq áp dụng cơ chế tập trung.

Chúng ta có thể sử dụng Transformer đã được huấn luyện để dịch một số câu đơn giản dưới đây.

for sentence in ['Go .', 'Wow !', "I'm OK .", 'I won !']:
    print(sentence + ' => ' + d2l.predict_s2s_ch9(
        model, sentence, src_vocab, tgt_vocab, num_steps, ctx))
Go . => entrez !
Wow ! => ça <unk> <unk> <unk> <unk> !
I'm OK . => ça va .
I won ! => j'ai gagné !

10.3.8. Tóm tắt

  • Mô hình Transformer dựa trên kiến ​​trúc mã hóa - giải mã.
  • Tầng tập trung đa đầu gồm có \(h\) tầng tập trung song song.
  • Mạng truyền xuôi theo vị trí gồm hai tầng kết nối đầy đủ được áp dụng trên chiều cuối cùng.
  • Chuẩn hóa theo tầng được áp dụng trên chiều cuối cùng (chiều đặc trưng), thay vì chiều đầu tiên (kích thước batch) như ở chuẩn hóa theo batch.
  • Biểu diễn vị trí là nơi duy nhất đưa thông tin vị trí trong chuỗi vào mô hình Transformer.

10.3.9. Bài tập

  1. Hãy thử huấn luyện với nhiều epoch hơn và so sánh mất mát giữa seq2seq và Transformer.
  2. Biểu diễn vị trí còn có lợi ích gì khác không?
  3. So sánh chuẩn hóa theo tầng với chuẩn hóa theo batch. Khi nào nên sử dụng một trong hai tầng chuẩn hóa thay vì tầng còn lại?

10.3.10. Thảo luận

10.3.11. Những người thực hiện

Bản dịch trong trang này được thực hiện bởi:

  • Đoàn Võ Duy Thanh
  • Nguyễn Duy Du
  • Lê Khắc Hồng Phúc
  • Trần Yến Thy
  • Phạm Minh Đức
  • Nguyễn Văn Quang
  • Phạm Hồng Vinh
  • Nguyễn Văn Cường
  • Nguyễn Cảnh Thướng