Khung toán học để giải mã Transformer: Khi ta thử "dịch ngược" trí tuệ nhân tạo

Các mô hình ngôn ngữ transformer như GPT-3 đang ngày càng được ứng dụng rộng rãi, nhưng càng mở rộng quy mô, chúng càng bộc lộ những hành vi khó lường và đôi khi gây hại. Một nhóm nghiên cứu đã đề xuất khung toán học để "dịch ngược" transformer thành các thành phần con người có thể hiểu được — một bước tiến quan trọng trong lĩnh vực diễn giải cơ chế. Đặc biệt, họ phát hiện ra induction head, một cơ chế then chốt giải thích khả năng học trong ngữ cảnh của mô hình.

Vấn đề: Mô hình càng lớn, càng khó hiểu

Các mô hình ngôn ngữ transformer ngày nay như GPT-3, LaMDA, Codex, Gopher đang được sử dụng ngày càng rộng rãi trong thực tế. Nhưng khi quy mô tăng lên, tính mở và dung lượng lớn của chúng tạo ra ngày càng nhiều khả năng xuất hiện những hành vi bất ngờ, thậm chí có hại.

Nhiều năm sau khi một mô hình lớn được huấn luyện, cả nhà phát triển lẫn người dùng vẫn thường xuyên phát hiện ra những khả năng — bao gồm cả hành vi có vấn đề — mà trước đó họ hoàn toàn không biết.

Một hướng tiếp cận để giải quyết vấn đề này là diễn giải cơ chế (mechanistic interpretability): cố gắng dịch ngược các phép tính chi tiết mà transformer thực hiện, giống như lập trình viên dịch ngược một tệp nhị phân phức tạp thành mã nguồn mà con người có thể đọc được. Nếu làm được điều này, ta có thể có 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, phát hiện vấn đề mới, và thậm chí dự đoán trước những rủi ro của các mô hình tương lai chưa được xây dựng.

Chiến lược: Bắt đầu từ mô hình đơn giản nhất

Với độ phức tạp khổng lồ của các mô hình ngôn ngữ hiện đại, nhóm nghiên cứu cho rằng cách hiệu quả nhất là bắt đầu từ những mô hình đơn giản nhất có thể rồi từ đó đi lên. Mục tiêu là tìm ra các mẫu thuật toán, mô-típ hoặc khung lý thuyết đơn giản, sau đó áp dụng cho các mô hình lớn và phức tạp hơn.

Cụ thể, bài báo nghiên cứu các transformer có từ hai lớp trở xuống và chỉ có khối attention — khác hoàn toàn với những mô hình lớn như GPT-3 với 96 lớp, xen kẽ giữa khối attention và khối MLP.

Bằng cách khái niệm hóa hoạt động của transformer theo một cách mới nhưng tương đương về mặt toán học, nhóm nghiên cứu có thể hiểu được các mô hình nhỏ này và nắm bắt đáng kể cách chúng vận hành bên trong.

Phát hiện quan trọng: Induction Head

Điểm đáng chú ý nhất là việc phát hiện ra các attention head đặc biệt mà nhóm nghiên cứu gọi là "induction head". Những head này có thể giải thích khả năng học trong ngữ cảnh (in-context learning) của các mô hình nhỏ, và chúng chỉ hình thành ở những mô hình có ít nhất hai lớp attention.

Nói cách khác, khi ta cho mô hình một chuỗi các token, induction head giúp mô hình nhận ra các mẫu lặp lại và dự đoán token tiếp theo dựa trên những gì đã xuất hiện trước đó — một dạng học tập rất cơ bản nhưng mạnh mẽ.

Cấu trúc mới để hiểu Transformer

Nhóm nghiên cứu đề xuất một cách nhìn mới về transformer thông qua khái niệm residual stream (dòng dư):

• Mỗi lớp không xử lý dữ liệu trực tiếp mà chỉ đọc từ residual stream và ghi kết quả trở lại
• Residual stream là một vector không gian nhiều chiều, hoạt động như một kênh truyền thông giữa các lớp
• Điều này có nghĩa là thông tin đã được ghi vào sẽ tồn tại mãi, trừ khi có lớp khác chủ động xóa nó
• Các lớp có thể gửi thông tin đến những lớp cụ thể khác thông qua các không gian con khác nhau

Một hệ quả thú vị của cấu trúc tuyến tính này là ta có thể nghĩ đến các "trọng số ảo" — kết nối trực tiếp giữa bất kỳ cặp lớp nào, kể cả khi chúng cách xa nhau nhiều lớp. Điều này giúp việc phân tích trở nên dễ dàng hơn.

Mô hình một lớp: Bigram và skip-trigram

Với các transformer một lớp chỉ có attention, nhóm nghiên cứu chứng minh chúng có thể được hiểu như tập hợp của một mô hình bigram và nhiều mô hình "skip-trigram".

Cụ thể, mỗi attention head có thể chọn lọc thông tin từ một token trước đó và sao chép thông tin để điều chỉnh xác suất của token tiếp theo. Các tương tác này tạo thành các skip-trigram có dạng [nguồn]...[đích][đầu ra].

Điều thú vị là khi mở rộng các ma trận QK và OV, ta có thể đọc trực tiếp các skip-trigram từ đó:

• Hầu hết các attention head dành phần lớn dung lượng cho việc sao chép (copying)
• Token nếu được chú ý đến sẽ làm tăng xác suất của chính nó xuất hiện
• Ví dụ như mô hình học được rằng "Barack" thường được theo sau bởi "Obama"
• Mô hình còn học được cấu trúc LaTeX và các chuỗi escape HTML

Đây có thể coi là một dạng học trong ngữ cảnh rất sơ khai, nhưng nó cho thấy tiềm năng của transformer trong việc thích ứng với ngữ cảnh.

Ý nghĩa đối với cộng đồng AI Việt Nam

Với các nhà nghiên cứu và kỹ sư AI tại Việt Nam đang làm việc với các mô hình ngôn ngữ, bài báo này mang lại một số bài học thiết thực:

• Hiểu rõ cơ chế bên trong giúp cải thiện chất lượng mô hình, đặc biệt khi tinh chỉnh (fine-tune) cho tiếng Việt
• Khả năng sao chép token của attention head giải thích tại sao mô hình có thể gặp lỗi với các từ tiếng Việt bị tách thành nhiều token
• Induction head là chìa khóa để hiểu tại sao mô hình có thể học từ vài ví dụ trong prompt (few-shot learning) — điều đặc biệt quan trọng khi ta muốn xây dựng ứng dụng mà không cần huấn luyện lại từ đầu
• Việc nhận diện "lỗi" trong mô hình (ví dụ như tăng xác suất của các cụm từ vô nghĩa) giúp chúng ta kiểm soát và cải thiện chất lượng đầu ra

Kết luận

Bài báo này là một bước đi tiên phong rất sơ khai trong việc dịch ngược transformer. Nhóm nghiên cứu thừa nhận còn rất xa mới có thể dịch ngược hoàn toàn các mô hình lớn, nhưng khung toán học và khái niệm induction head đã chứng minh được tính hữu ích.

Chúng tôi đã chuyển vấn đề hiểu transformer một lớp chỉ có attention thành vấn đề hiểu các ma trận OV và QK mở rộng của chúng.

Dù các ma trận này khổng lồ (với hàng tỷ phần tử), việc phân tích chúng đã mở ra cánh cửa để hiểu sâu hơn về cách mà những cỗ máy ngôn ngữ mạnh mẽ nhất hiện nay thực sự vận hành — một bước tiến quan trọng hướng tới AI an toàn và có thể kiểm chứng được.