Trực quan hóa cơ chế Attention của LLM: Hiểu sâu hơn cách mô hình sao chép và suy luận

Trong thế giới trí tuệ nhân tạo, các mô hình ngôn ngữ lớn (LLM) dựa trên kiến trúc transformer luôn được nhắc đến như một "hộp đen" huyền bí. Chúng ta biết chúng hoạt động, biết chúng thông minh, nhưng hiểu sâu về cách chúng liên kết các thông tin lại là một câu chuyện khác.

Một công cụ trực quan hóa mới vừa xuất hiện trên Hacker News, được xây dựng bởi nhà phát triển Isham F., hứa hẹn sẽ giúp chúng ta "nhìn xuyên" vào bên trong chiếc hộp đen đó, khám phá cách mô hình quyết định token nào sẽ ảnh hưởng đến token tiếp theo.

Cơ chế "vẽ lại" quá trình suy luận

Ý tưởng của công cụ này rất độc đáo và trực quan. Khi mô hình tạo văn bản, người dùng có thể nhấp hoặc di chuột qua từng token được tạo để xem những token nào trong quá khứ đã "tác động" đến quá trình tạo ra token đó. Sự tác động này được thể hiện qua độ mờ (opacity) của các token trước đó - token càng ảnh hưởng nhiều thì càng được tô đậm.

Điểm thú vị là tác giả không dùng phép đo phức tạp nào mà chỉ đơn giản là tính toán trọng số attention, nhân với độ lớn của vector giá trị (value vector), tổng hợp qua tất cả các đầu attention và các lớp trong mạng. Kết quả sau đó được chuẩn hóa để hiển thị một con số duy nhất cho mỗi token trong quá khứ.

Giải mã khả năng "sao chép" và "tổng hợp" thông tin

Một trong những phát hiện thú vị nhất mà công cụ này mang lại là làm sáng tỏ khả năng sao chép thông tin gần như chính xác tuyệt đối của LLM - một điều mà nhiều người (kể cả tác giả) từng thấy khó hiểu.

Nếu LLM hoạt động bằng cách dự đoán token tiếp theo một cách xác suất, tại sao chúng lại sao chép chính xác đến vậy? Lẽ ra chúng phải gây ra lỗi ngẫu nhiên nào đó chứ?

Câu trả lời nằm ở cơ chế attention. LLM không dự đoán toàn bộ chuỗi thông tin từ một trạng thái nội bộ hạn chế. Thay vào đó, vì có quyền truy cập vào toàn bộ token quá khứ, chúng có thể "chọn lọc" thông tin cần thiết để sao chép một cách chính xác nhất. Trong ví dụ về "Debugging an Average Function" của tác giả, một mô hình chỉ với 600 triệu tham số (tương đối nhỏ) vẫn có thể tái tạo hoàn hảo một hàm JavaScript ngoài phần chỉnh sửa theo yêu cầu.

Ngoài khả năng sao chép, công cụ còn cho thấy cách mô hình kết hợp thông tin từ nhiều nguồn khác nhau. Ví dụ, khi di chuột qua từ "remain" trong câu "Existing access cards and phone numbers remain" (từ mẫu Office Move Summary), người dùng có thể thấy mô hình lấy thông tin từ cả "work" trong "Existing employee access cards will work" lẫn "stay the same" trong "company phone numbers will stay the same". Điều này cho thấy LLM có thể tổng hợp một khái niệm mới từ việc kết hợp ý nghĩa của nhiều cụm từ khác nhau - một cách suy luận khá tinh vi.

Những thách thức kỹ thuật trong quá trình phát triển

Về mặt kỹ thuật, việc xây dựng công cụ này không hề đơn giản. Tác giả sử dụng React và Transformers.js để tạo ứng dụng web. Tuy nhiên, điểm mấu chốt là ứng dụng cần truy cập vào dữ liệu nội bộ của mô hình trong quá trình tạo văn bản. Vòng lặp tạo sinh chuẩn không cho phép điều này, vì vậy tác giả đã phải "vibe-code" một vòng lặp tùy chỉnh để theo dõi các giá trị attention trong thời gian thực.

Thách thức lớn nhất là các file mô hình .onnx được sử dụng bởi Transformers.js không cho phép truy cập vào các đầu ra trung gian bên trong. Để giải quyết, tác giả đã viết một script nhỏ để sửa đổi file mô hình, "phơi bày" các giá trị nội bộ ra ngoài như một đầu ra phụ. Kết quả là phải tạo một phiên bản mô hình đặc biệt, tải lên Hugging Face và cấu hình ứng dụng trỏ đến đó.

Lời kết

Công cụ này không chỉ là một sản phẩm kỹ thuật thú vị mà còn là một công cụ học tập tuyệt vời cho những ai muốn hiểu sâu về cách LLM hoạt động. Với việc đã tạo sẵn các prompt mẫu hiển thị tức thì, người dùng không cần phải tải mô hình lớn về máy mà có thể trải nghiệm ngay trên trình duyệt tại đây. Mã nguồn cũng đã được công khai trên GitHub cho những ai muốn phát triển thêm.