Khung toán học cho mạch Transformer (2021): Nền tảng của khả năng diễn giải AI

Mục lục

Các mô hình ngôn ngữ Transformer là một công nghệ mới nổi đang ngày càng được sử dụng rộng rãi trong thực tế, ví dụ như trong các hệ thống như GPT-3, LaMDA, Codex, Meena, Gopher và các mô hình tương tự. Tuy nhiên, khi các mô hình này được mở rộng quy mô, tính chất mở và năng lực cao của chúng tạo ra phạm vi ngày càng lớn cho các hành vi bất ngờ và đôi khi có hại. Ngay cả nhiều năm sau khi một mô hình lớn được huấn luyện, cả người tạo ra và người dùng vẫn thường xuyên phát hiện ra các khả năng của mô hình – bao gồm cả các hành vi có vấn đề – mà trước đây họ không hề hay biết.

Một hướng giải quyết các vấn đề này là khả năng diễn giải cơ học (mechanistic interpretability), cố gắng kỹ thuật đảo ngược (reverse engineer) các tính toán chi tiết được thực hiện bởi các Transformer, tương tự như cách một lập trình viên có thể cố gắng kỹ thuật đảo ngược các tệp nhị phân phức tạp thành mã nguồn mà con người có thể đọc được. Nếu điều này khả thi, nó có khả năng cung cấp một cách tiếp cận hệ thống hơn để giải thích các vấn đề an toàn hiện tại, xác định các vấn đề mới và thậm chí có thể dự đoán các vấn đề an toàn của các mô hình mạnh mẽ trong tương lai chưa được xây dựng. Một dự án trước đây, chuỗi Distill Circuits, đã cố gắng kỹ thuật đảo ngược các mô hình thị giác, nhưng cho đến nay vẫn chưa có dự án tương đương cho các Transformer hoặc mô hình ngôn ngữ.

Trong bài báo này, chúng tôi cố gắng thực hiện những bước đầu tiên, rất sơ khai hướng tới việc kỹ thuật đảo ngược các Transformer. Với độ phức tạp và quy mô đáng kinh ngạc của các mô hình ngôn ngữ hiện đại, chúng tôi nhận thấy việc bắt đầu với các mô hình đơn giản nhất có thể và phát triển dần từ đó là hiệu quả nhất. Mục tiêu của chúng tôi là khám phá các mô hình thuật toán, các mô típ hoặc khung làm việc đơn giản mà sau đó có thể áp dụng cho các mô hình lớn hơn và phức tạp hơn. Cụ thể, trong bài báo này, chúng tôi sẽ nghiên cứu các Transformer có từ hai lớp trở xuống chỉ có các khối attention – điều này trái ngược với một Transformer hiện đại, lớn như GPT-3, có 96 lớp và xen kẽ các khối attention với các khối MLP.

Chúng tôi nhận thấy rằng bằng cách khái niệm hóa hoạt động của các Transformer theo một cách mới nhưng tương đương về mặt toán học, chúng tôi có thể hiểu được các mô hình nhỏ này và đạt được sự hiểu biết đáng kể về cách chúng vận hành bên trong. Đáng chú ý, chúng tôi phát hiện ra rằng các đầu attention (attention heads) cụ thể mà chúng tôi gọi là "induction heads" có thể giải thích việc học trong ngữ cảnh (in-context learning) ở các mô hình nhỏ này, và các đầu này chỉ phát triển trong các mô hình có ít nhất hai lớp attention. Chúng tôi cũng xem xét một số ví dụ về các đầu này hoạt động trên dữ liệu cụ thể.

Chúng tôi không cố gắng áp dụng những hiểu biết của mình vào các mô hình lớn hơn trong bài báo đầu tiên này, nhưng trong một bài báo sắp tới, chúng tôi sẽ chỉ ra rằng cả khung toán học của chúng tôi để hiểu các Transformer và khái niệm về induction heads đều tiếp tục có liên quan ít nhất một phần đối với các mô hình lớn hơn và thực tế hơn nhiều – mặc dù chúng tôi vẫn còn một chặng đường rất dài để có thể kỹ thuật đảo ngược hoàn toàn các mô hình như vậy.

Tóm tắt kết quả

Kết quả kỹ thuật đảo ngược

Để khám phá thách thức của việc kỹ thuật đảo ngược các Transformer, chúng tôi đã kỹ thuật đảo ngược một số mô hình đồ chơi (toy models) chỉ sử dụng attention. Khi làm như vậy, chúng tôi nhận thấy:

Các bài học khái niệm

Chúng tôi nhận thấy rằng nhiều chi tiết tinh tế của kiến trúc Transformer đòi hỏi chúng tôi phải tiếp cận việc kỹ thuật đảo ngược nó theo một cách khá khác biệt so với cách hoạt động của InceptionV1 Circuits. Chúng tôi sẽ giải mã từng điểm này trong các phần dưới đây, nhưng hiện tại chúng tôi tóm tắt ngắn gọn. Chúng tôi cũng sẽ mở rộng nhiều thuật ngữ mà chúng tôi giới thiệu ở đây khi đến các phần thích hợp. (Để rõ ràng, chúng tôi không có ý định khẳng định rằng bất kỳ điểm nào trong số này nhất thiết phải là mới; nhiều điểm đã hiện diện một cách ngầm định hoặc rõ ràng trong các bài báo khác.)

Tổng quan về Transformer

Trước khi chúng tôi cố gắng kỹ thuật đảo ngược các Transformer, việc xem xét ngắn gọn cấu trúc cấp cao của Transformer và mô tả cách chúng tôi suy nghĩ về chúng là rất hữu ích.

Trong nhiều trường hợp, chúng tôi thấy việc định hình lại các Transformer theo những cách tương đương nhưng không chuẩn là rất hữu ích. Khả năng diễn giải cơ học đòi hỏi chúng tôi phải chia nhỏ các mô hình thành các phần mà con người có thể diễn giải được. Một bước đầu tiên quan trọng là tìm ra cách biểu diễn giúp việc suy luận về mô hình trở nên dễ dàng nhất. Trong học sâu hiện đại, có rất nhiều sự nhấn mạnh vào hiệu quả tính toán – và điều này là có lý do! – và các mô tả toán học của chúng tôi về các mô hình thường phản ánh các quyết định về cách viết mã hiệu quả để chạy mô hình. Nhưng khi có nhiều cách tương đương để biểu diễn cùng một phép tính, khả năng cao là cách biểu diễn dễ hiểu nhất đối với con người và cách biểu diễn hiệu quả nhất về mặt tính toán sẽ khác nhau.

Việc xem xét các Transformer cũng sẽ cho phép chúng tôi thống nhất về thuật ngữ, vốn đôi khi có thể thay đổi. Chúng tôi cũng sẽ giới thiệu một số ký hiệu trong quá trình này, nhưng vì các ký hiệu này được sử dụng trong nhiều phần, chúng tôi cung cấp mô tả chi tiết về tất cả các ký hiệu trong phụ lục ký hiệu như một tài liệu tham khảo ngắn gọn cho độc giả.

Đơn giản hóa mô hình

Để chứng minh các ý tưởng trong bài báo này ở dạng rõ ràng nhất, chúng tôi tập trung vào các "toy transformers" với một số đơn giản hóa.

Trong hầu hết các phần của bài báo này, chúng tôi sẽ thực hiện một thay đổi rất đáng kể: chúng tôi tập trung vào các Transformer "chỉ có attention", không có các lớp MLP. Đây là một sự đơn giản hóa rất mạnh mẽ đối với kiến trúc Transformer. Chúng tôi được thúc đẩy một phần bởi thực tế là các mạch với các đầu attention đặt ra những thách thức mới mà công trình Distill circuits chưa gặp phải, và việc xem xét chúng một cách cô lập cho phép chúng tôi xử lý các vấn đề đó một cách đặc biệt thanh lịch. Nhưng chúng tôi cũng đơn giản là đã ít thành công hơn nhiều trong việc hiểu các lớp MLP cho đến nay; trong các Transformer bình thường có cả lớp attention và lớp MLP, có nhiều mạch được điều phối chủ yếu bởi các đầu attention mà chúng tôi có thể nghiên cứu, một số trong đó có vẻ rất quan trọng, nhưng các phần MLP lại khó nắm bắt hơn nhiều. Đây là một điểm yếu lớn trong công trình của chúng tôi mà chúng tôi dự định tập trung giải quyết trong tương lai. Mặc dù vậy, chúng tôi sẽ có một số thảo luận về các Transformer có lớp MLP trong các phần sau.

Chúng tôi cũng thực hiện một số thay đổi mà chúng tôi coi là bề ngoài hơn và chủ yếu được thực hiện để rõ ràng và đơn giản. Chúng tôi không xem xét các độ lệch (biases), nhưng một mô hình có độ lệch luôn có thể được mô phỏng mà không cần chúng bằng cách gộp chúng vào các trọng số và tạo ra một chiều luôn bằng một. Ngoài ra, các độ lệch trong các Transformer chỉ có attention chủ yếu nhân lên để trở thành các độ lệch trên các logit. Chúng tôi cũng bỏ qua chuẩn hóa lớp (layer normalization). Việc xem xét rõ ràng nó làm tăng thêm khá nhiều độ phức tạp, và cho đến một tỷ lệ biến đổi, layer norm có thể được gộp vào các trọng số liền kề. Chúng tôi cũng hy vọng rằng, trừ một số phiền toái về triển khai, layer norm có thể được thay thế bằng batch normalization (có thể được gộp hoàn toàn vào các tham số liền kề).

Kiến trúc cấp cao

Có một số biến thể của các mô hình ngôn ngữ Transformer. Chúng tôi tập trung vào các mô hình ngôn ngữ Transformer tự hồi quy, chỉ có bộ giải mã (decoder-only), chẳng hạn như GPT-3. (Bài báo Transformer gốc có cấu trúc bộ mã hóa-giải mã đặc biệt để hỗ trợ dịch thuật, nhưng nhiều mô hình ngôn ngữ hiện đại không bao gồm cấu trúc này.)

Một Transformer bắt đầu bằng một token embedding, theo sau là một loạt các "khối dư" (residual blocks), và cuối cùng là một token unembedding. Mỗi khối dư bao gồm một lớp attention, theo sau là một lớp MLP. Cả lớp attention và lớp MLP đều "đọc" đầu vào của chúng từ luồng dư (residual stream) (bằng cách thực hiện một phép chiếu tuyến tính), và sau đó "ghi" kết quả của chúng vào luồng dư bằng cách cộng thêm một phép chiếu tuyến tính trở lại. Mỗi lớp attention bao gồm nhiều đầu (heads), hoạt động song song.

Trọng số ảo và luồng dư như một kênh truyền thông

Trọng số ảo và luồng dư như một kênh truyền thông

Một trong những đặc điểm chính của kiến trúc cấp cao của một Transformer là mỗi lớp cộng kết quả của nó vào cái mà chúng tôi gọi là "luồng dư" (residual stream). Việc xây dựng các mô hình với luồng dư bắt nguồn từ các công trình ban đầu của nhóm Schmidhuber, chẳng hạn như mạng đường cao tốc (highway networks) và LSTM, vốn đã đạt được thành công hiện đại đáng kể trong kiến trúc mạng dư (residual network) gần đây hơn. Trong các Transformer, các vectơ luồng dư thường được gọi là "embedding". Chúng tôi thích thuật ngữ luồng dư hơn, vì nó nhấn mạnh bản chất dư (mà chúng tôi tin là quan trọng) và cũng vì chúng tôi tin rằng luồng dư thường dành các không gian con cho các token khác ngoài token hiện tại, phá vỡ các trực giác mà thuật ngữ embedding gợi ý. Luồng dư đơn giản là tổng đầu ra của tất cả các lớp trước đó và embedding ban đầu. Chúng tôi thường coi luồng dư như một kênh truyền thông, vì bản thân nó không thực hiện bất kỳ quá trình xử lý nào và tất cả các lớp đều giao tiếp thông qua nó.

Luồng dư có cấu trúc tuyến tính sâu sắc. Cần lưu ý rằng luồng dư hoàn toàn tuyến tính là rất bất thường trong các kiến trúc mạng thần kinh: ngay cả ResNet, kiến trúc tương tự nhất được sử dụng rộng rãi, cũng có các hàm kích hoạt phi tuyến tính trên luồng dư của chúng, hoặc được áp dụng bất cứ khi nào luồng dư được truy cập! Mỗi lớp thực hiện một phép biến đổi tuyến tính tùy ý để "đọc vào" thông tin từ luồng dư ở phần đầu, và thực hiện một phép biến đổi tuyến tính tùy ý khác trước khi cộng để "ghi" đầu ra của nó trở lại luồng dư. Cấu trúc cộng tuyến tính này của luồng dư có rất nhiều ý nghĩa quan trọng. Một hệ quả cơ bản là luồng dư không có "cơ sở đặc quyền"; chúng ta có thể xoay nó bằng cách xoay tất cả các ma trận tương tác với nó mà không làm thay đổi hành vi của mô hình.

Trọng số ảo

Một hệ quả đặc biệt hữu ích của việc luồng dư là tuyến tính là người ta có thể nghĩ về các "trọng số ảo" ngầm định kết nối trực tiếp bất kỳ cặp lớp nào (ngay cả những lớp cách nhau bởi nhiều lớp khác), bằng cách nhân các tương tác của chúng thông qua luồng dư. Các trọng số ảo này là tích của các trọng số đầu ra của một lớp với các trọng số đầu vào của lớp khác (ví dụ: W_{I}^2W_{O}^1), và mô tả mức độ mà một lớp sau đọc thông tin được ghi bởi một lớp trước đó.

Không gian con và băng thông luồng dư

Luồng dư là một không gian vectơ nhiều chiều. Trong các mô hình nhỏ, nó có thể có hàng trăm chiều; trong các mô hình lớn, nó có thể lên tới hàng chục nghìn. Điều này có nghĩa là các lớp có thể gửi thông tin khác nhau đến các lớp khác nhau bằng cách lưu trữ nó trong các không gian con khác nhau. Điều này đặc biệt quan trọng trong trường hợp các đầu attention, vì mỗi đầu riêng lẻ hoạt động trên các không gian con tương đối nhỏ (thường là 64 hoặc 128 chiều), và có thể rất dễ dàng ghi vào các không gian con hoàn toàn tách biệt và không tương tác.

Sau khi được thêm vào, thông tin sẽ tồn tại trong một không gian con trừ khi một lớp khác chủ động xóa nó. Từ góc độ này, các chiều của luồng dư trở thành thứ gì đó giống như "bộ nhớ" hoặc "băng thông". Các token embedding ban đầu, cũng như các unembedding, chủ yếu tương tác với một phần tương đối nhỏ của các chiều. Điều này để lại hầu hết các chiều "trống" cho các lớp khác lưu trữ thông tin.

Có vẻ như chúng ta nên mong đợi băng thông luồng dư sẽ có nhu cầu rất cao! Nhìn chung, có nhiều "chiều tính toán" (chẳng hạn như các nơ-ron và các chiều kết quả của đầu attention) hơn nhiều so với số chiều mà luồng dư có để di chuyển thông tin. Chỉ một lớp MLP điển hình đã có số nơ-ron gấp bốn lần số chiều của luồng dư. Vì vậy, ví dụ, tại lớp 25 của một Transformer 50 lớp, luồng dư có số nơ-ron gấp 100 lần số chiều trước đó, cố gắng giao tiếp với số nơ-ron gấp 100 lần số chiều sau đó, bằng cách nào đó giao tiếp trong sự chồng chất! Chúng tôi gọi các tensor như thế này là "kích hoạt nút thắt" (bottleneck activations) và hy vọng chúng sẽ là những thách thức bất thường để diễn giải. (Đây là lý do chính tại sao chúng tôi sẽ cố gắng tách các luồng giao tiếp khác nhau đang diễn ra thông qua luồng dư theo các trọng số ảo, thay vì nghiên cứu trực tiếp nó.)

Có lẽ vì nhu cầu cao này đối với băng thông luồng dư, chúng tôi đã thấy những gợi ý rằng một số nơ-ron MLP và đầu attention có thể thực hiện một loại vai trò "quản lý bộ nhớ", xóa các chiều luồng dư được thiết lập bởi các lớp khác bằng cách đọc thông tin và ghi ra phiên bản phủ định. Cần lưu ý rằng mặc dù đây có thể là các cơ chế chung để xóa thông tin "quản lý bộ nhớ", chúng cũng có thể là các cơ chế để xóa thông tin có điều kiện, chỉ hoạt động trong một số trường hợp.

Các đầu Attention là độc lập và có tính cộng

Các đầu Attention là độc lập và có tính cộng

Như đã thấy ở trên, chúng tôi coi các lớp attention của Transformer là một vài đầu attention hoàn toàn độc lập h\in H hoạt động hoàn toàn song song và mỗi đầu cộng đầu ra của nó trở lại luồng dư. Nhưng đây không phải là cách các lớp Transformer thường được trình bày, và có thể không rõ ràng rằng chúng tương đương nhau.

Trong bài báo gốc của Vaswani và cộng sự về Transformer, đầu ra của một lớp attention được mô tả bằng cách xếp chồng các vectơ kết quả r^{h_1}, r^{h_2},..., và sau đó nhân với một ma trận đầu ra W_O^H. Hãy chia W_O^H thành các khối có kích thước bằng nhau cho mỗi đầu [W_O^{h_1}, W_O^{h_2}...]. Sau đó, chúng tôi quan sát thấy rằng:

W_O^H \left[\begin{matrix}r^{h_1}\\r^{h_2}\\... \end{matrix}\right] ~~=~~ \left[W_O^{h_1},~ W_O^{h_2},~... \right]\cdot\left[\begin{matrix}r^{h_1}\\r^{h_2}\\...\end{matrix}\right] ~~=~~ \sum_i W_O^{h_i} r^{h_i}

Tiết lộ rằng nó tương đương với việc chạy các đầu một cách độc lập, nhân mỗi đầu với ma trận đầu ra riêng của nó và cộng chúng vào luồng dư. Định nghĩa nối (concatenate) thường được ưu tiên vì nó tạo ra một phép nhân ma trận lớn hơn và hiệu quả tính toán hơn. Nhưng để hiểu các Transformer về mặt lý thuyết, chúng tôi thích coi chúng là có tính cộng độc lập.

Các đầu Attention như là sự di chuyển thông tin