10.1. Cơ chế Tập trung¶

Trong Section 9.7, chúng ta dùng mạng hồi tiếp để mã hóa thông tin của chuỗi nguồn đầu vào thành trạng thái ẩn và truyền nó tới bộ giải mã để sinh chuỗi đích. Một token trong chuỗi đích có thể chỉ liên quan mật thiết tới một vài token chứ không nhất thiết là toàn bộ token trong chuỗi nguồn. Ví dụ, khi dịch “Hello world.” thành “Bonjour le monde.”, từ “Bonjour” ánh xạ tới từ “Hello” và từ “monde” ánh xạ tới từ “world”. Trong mô hình seq2seq, bộ giải mã có thể ngầm chọn thông tin tương ứng từ trạng thái ẩn được truyền đến từ bộ mã hóa. Tuy nhiên, cơ chế tập trung (attention mechanism) thực hiện phép chọn này một cách tường minh.

Cơ chế tập trung có thể được coi là phép gộp tổng quát. Nó gộp đầu vào dựa trên các trọng số khác nhau. Thành phần cốt lõi của cơ chế tập trung là tầng tập trung. Đầu vào của tầng tập trung được gọi ngắn gọn là câu truy vấn (query). Với mỗi câu truy vấn, tầng tập trung trả về đầu ra dựa trên bộ nhớ là tập các cặp khóa-giá trị được mã hóa trong tầng tập trung này. Cụ thể, giả sử bộ nhớ chứa \(n\) cặp vector khóa-giá trị, \((\mathbf{k}_1, \mathbf{v}_1), \ldots, (\mathbf{k}_n, \mathbf{v}_n)\), với \(\mathbf{k}_i \in \mathbb R^{d_k}\), \(\mathbf{v}_i \in \mathbb R^{d_v}\). Với mỗi vector truy vấn \(\mathbf{q} \in \mathbb R^{d_q}\), tầng tập trung trả về đầu ra \(\mathbf{o} \in \mathbb R^{d_v}\) có cùng kích thước với vector giá trị.

Fig. 10.1.1 Tầng tập trung trả về giá trị dựa trên câu truy vấn đầu vào và bộ nhớ của nó.¶

Chi tiết về cơ chế tập trung được minh họa trong Fig. 10.1.2. Để tính toán đầu ra của tầng tập trung, chúng ta sử dụng hàm tính điểm \(\alpha\) để đo độ tương đồng giữa câu truy vấn và các khóa. Sau đó, với mỗi khóa \((\mathbf{k}_1, \mathbf{v}_1), \ldots, (\mathbf{k}_n, \mathbf{v}_n)\), ta tính điểm \(a_1, \ldots, a_n\) như sau:

(10.1.1)¶\[a_i = \alpha(\mathbf q, \mathbf k_i).\]

Tiếp theo, chúng ta sử dụng hàm softmax để thu được các trọng số tập trung (attention weights), cụ thể:

(10.1.2)¶\[\mathbf{b} = \mathrm{softmax}(\mathbf{a})\quad \text{trong đó }\quad {b}_i = \frac{\exp(a_i)}{\sum_j \exp(a_j)}, \mathbf{b} = [b_1, \ldots, b_n]^T .\]

Cuối cùng, đầu ra của tầng là tổng trọng số của các giá trị:

(10.1.3)¶\[\mathbf o = \sum_{i=1}^n b_i \mathbf v_i.\]

Fig. 10.1.2 Đầu ra của tầng tập trung là tổng trọng số của các giá trị.¶

Cách lựa chọn hàm tính điểm khác nhau sẽ tạo ra các tầng tập trung khác nhau. Ở dưới đây chúng tôi sẽ trình bày hai tầng tập trung thường hay được sử dụng. Đầu tiên chúng tôi giới thiệu hai toán tử cần thiết để lập trình hai tầng này: toán tử softmax có mặt nạ masked_softmax và toán tử tích vô hướng chuyên biệt theo batch batched_dot.

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

Toán tử softmax có mặt nạ nhận đầu vào là một tensor 3 chiều và cho phép ta lọc ra một số phần tử bằng cách xác định độ dài hợp lệ cho chiều cuối cùng. (Tham khảo Section 9.5 về định nghĩa của độ dài hợp lệ). Do đó, những giá trị nằm ngoài độ dài hợp lệ sẽ được gán bằng \(0\). Chúng ta lập trình hàm masked_softmax như sau.

# Saved in the d2l package for later use
def masked_softmax(X, valid_len):
    # X: 3-D tensor, valid_len: 1-D or 2-D tensor
    if valid_len is None:
        return npx.softmax(X)
    else:
        shape = X.shape
        if valid_len.ndim == 1:
            valid_len = valid_len.repeat(shape[1], axis=0)
        else:
            valid_len = valid_len.reshape(-1)
        # Fill masked elements with a large negative, whose exp is 0
        X = npx.sequence_mask(X.reshape(-1, shape[-1]), valid_len, True,
                              axis=1, value=-1e6)
        return npx.softmax(X).reshape(shape)

Để minh họa cách hàm trên hoạt động, chúng ta hãy khởi tạo hai ma trận đầu vào kích thước là \(2 \times 4\). Bên cạnh đó, chúng ta sẽ gán độ dài hợp lệ cho mẫu thứ nhất là 2 và mẫu thứ hai là 3. Từ đó, những giá trị đầu ra nằm ngoài độ dài hợp lệ sẽ được gán bằng \(0\) như dưới đây.

masked_softmax(np.random.uniform(size=(2, 2, 4)), np.array([2, 3]))

array([[[0.488994  , 0.511006  , 0.        , 0.        ],
        [0.43654838, 0.56345165, 0.        , 0.        ]],

       [[0.28817102, 0.3519408 , 0.3598882 , 0.        ],
        [0.29034293, 0.25239873, 0.45725834, 0.        ]]])

Ngoài ra, toán tử thứ hai batched_dot nhận hai đầu vào là \(X\) và \(Y\) có kích thước lần lượt là \((b, n, m)\) và \((b, m, k)\), và trả về đầu ra có kích thước là \((b, n, k)\). Cụ thể, toán tử này tính \(b\) tích vô hướng với \(i= \{1,\ldots, b\}\) như sau:

(10.1.4)¶\[Z[i,:,:] = X[i,:,:] Y[i,:,:].\]

npx.batch_dot(np.ones((2, 1, 3)), np.ones((2, 3, 2)))

array([[[3., 3.]],

       [[3., 3.]]])

10.1.1. Tầng Tập trung Tích Vô hướng¶

Với hai toán tử masked_softmax và batched_dot ở trên, chúng ta sẽ đi vào chi tiết hai loại tầng tập trung được sử dụng phổ biến. Loại đầu tiên là tập trung tích vô hướng (dot product attention): nó giả định rằng câu truy vấn có cùng kích thước chiều với khóa, cụ thể là \(\mathbf q, \mathbf k_i \in\mathbb R^d\) với mọi \(i\). Tầng tập trung tích vô hướng sẽ tính điểm bằng cách lấy tích vô hướng giữa câu truy vấn và khóa, sau đó chia cho \(\sqrt{d}\) để giảm thiểu ảnh hưởng không liên quan của số chiều \(d\) lên điểm số. Nói cách khác,

(10.1.5)¶\[\alpha(\mathbf q, \mathbf k) = \langle \mathbf q, \mathbf k \rangle /\sqrt{d}.\]

Mở rộng ra từ các câu truy vấn và khóa một chiều, chúng ta luôn có thể tổng quát hóa chúng lên thành các giá trị truy vấn và khóa đa chiều. Giả định rằng \(\mathbf Q\in\mathbb R^{m\times d}\) chứa \(m\) câu truy vấn và \(\mathbf K\in\mathbb R^{n\times d}\) chứa toàn bộ \(n\) khóa. Chúng ta có thể tính toàn bộ \(mn\) điểm số như sau

(10.1.6)¶\[\alpha(\mathbf Q, \mathbf K) = \mathbf Q \mathbf K^\top /\sqrt{d}.\]

Với (10.1.6), chúng ta có thể lập trình tầng tập trung tích vô hướng DotProductAttention hỗ trợ một batch các câu truy vấn và các cặp khóa-giá trị. Ngoài ra, chúng ta cũng dùng thêm một tầng dropout để điều chuẩn.

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

    # query: (batch_size, #queries, d)
    # key: (batch_size, #kv_pairs, d)
    # value: (batch_size, #kv_pairs, dim_v)
    # valid_len: either (batch_size, ) or (batch_size, xx)
    def forward(self, query, key, value, valid_len=None):
        d = query.shape[-1]
        # Set transpose_b=True to swap the last two dimensions of key
        scores = npx.batch_dot(query, key, transpose_b=True) / math.sqrt(d)
        attention_weights = self.dropout(masked_softmax(scores, valid_len))
        return npx.batch_dot(attention_weights, value)

Hãy kiểm tra lớp DotProductAttention với một ví dụ nhỏ sau. Đầu tiên ta tạo 2 batch, mỗi batch có 1 câu truy vấn và 10 cặp khóa-giá trị. Thông qua đối số valid_len, ta chỉ định rằng ta sẽ kiểm tra \(2\) cặp khóa-giá trị đầu tiên cho batch đầu tiên và \(6\) cặp cho batch thứ hai. Do đó, mặc dù cả hai batch đều có cùng câu truy vấn và các cặp khóa-giá trị, chúng ta sẽ thu được các đầu ra khác nhau.

atten = DotProductAttention(dropout=0.5)
atten.initialize()
keys = np.ones((2, 10, 2))
values = np.arange(40).reshape(1, 10, 4).repeat(2, axis=0)
atten(np.ones((2, 1, 2)), keys, values, np.array([2, 6]))

array([[[ 2.      ,  3.      ,  4.      ,  5.      ]],

       [[10.      , 11.      , 12.000001, 13.      ]]])

Như đã thấy ở trên, tập trung tích vô hướng chỉ đơn thuần nhân câu truy vấn và khóa lại với nhau, hi vọng rằng từ đó thu được những điểm tương đồng giữa chúng. Tuy nhiên, câu truy vấn và khóa có thể không có cùng kích thước chiều. Để giải quyết vấn đề này, chúng ta cần nhờ đến cơ chế tập trung perceptron đa tầng.

10.1.2. Tập trung Perceptron Đa tầng¶

Trong cơ chế tập trung perceptron đa tầng (multilayer perceptron attention), chúng ta chiếu cả câu truy vấn và các khóa lên \(\mathbb R^{h}\) bằng các tham số trọng số được học. Giả định rằng các trọng số được học là \(\mathbf W_k\in\mathbb R^{h\times d_k}\), \(\mathbf W_q\in\mathbb R^{h\times d_q}\) và \(\mathbf v\in\mathbb R^{h}\). Hàm tính điểm sẽ được định nghĩa như sau

(10.1.7)¶\[\alpha(\mathbf k, \mathbf q) = \mathbf v^\top \text{tanh}(\mathbf W_k \mathbf k + \mathbf W_q\mathbf q).\]

Một cách trực quan, ta có thể tưởng tượng \(\mathbf W_k \mathbf k + \mathbf W_q\mathbf q\) chính là việc nối khóa và giá trị lại với nhau theo chiều đặc trưng và đưa chúng qua perceptron có một tầng ẩn với kích thước là \(h\) và tầng đầu ra với kích thước là \(1\). Trong tầng ẩn này, hàm kích hoạt là \(tanh\) và không có hệ số điều chỉnh. Giờ hãy lập trình một tầng tập trung perceptron đa tầng.

# Saved in the d2l package for later use
class MLPAttention(nn.Block):
    def __init__(self, units, dropout, **kwargs):
        super(MLPAttention, self).__init__(**kwargs)
        # Use flatten=True to keep query's and key's 3-D shapes
        self.W_k = nn.Dense(units, activation='tanh',
                            use_bias=False, flatten=False)
        self.W_q = nn.Dense(units, activation='tanh',
                            use_bias=False, flatten=False)
        self.v = nn.Dense(1, use_bias=False, flatten=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, valid_len):
        query, key = self.W_q(query), self.W_k(key)
        # Expand query to (batch_size, #querys, 1, units), and key to
        # (batch_size, 1, #kv_pairs, units). Then plus them with broadcast
        features = np.expand_dims(query, axis=2) + np.expand_dims(key, axis=1)
        scores = np.squeeze(self.v(features), axis=-1)
        attention_weights = self.dropout(masked_softmax(scores, valid_len))
        return npx.batch_dot(attention_weights, value)

Để kiểm tra lớp MLPAttention phía trên, chúng ta sẽ sử dụng lại đầu vào ở ví dụ đơn giản trước. Như ta thấy ở dưới, mặc dù MLPAttention chứa thêm một mô hình MLP, chúng ta vẫn thu được đầu ra tương tự DotProductAttention.

atten = MLPAttention(units=8, dropout=0.1)
atten.initialize()
atten(np.ones((2, 1, 2)), keys, values, np.array([2, 6]))

array([[[ 2.      ,  3.      ,  4.      ,  5.      ]],

       [[10.      , 11.      , 12.000001, 13.      ]]])

10.1.3. Tóm tắt¶

Tầng tập trung lựa chọn một cách tường minh các thông tin liên quan.
Ô nhớ của tầng tập trung chứa các cặp khóa-giá trị, do đó đầu ra của nó gần các giá trị có khóa giống với câu truy vấn.
Hai mô hình tập trung được sử dụng phổ biến là tập trung tích vô hướng và tập trung perceptron đa tầng.

10.1.4. Bài tập¶

Ưu và khuyết điểm của tầng tập trung tích vô hướng và tập trung perceptron đa tầng là gì?

10.1.5. Thảo luận¶

10.1.6. 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 Văn Quang
Nguyễn Cảnh Thướng
Nguyễn Văn Cường
Võ Tấn Phát
Lê Khắc Hồng Phúc
Phạm Minh Đức
Phạm Hồng Vinh