Show HN: Trực quan hóa cơ chế Attention của LLM ngay trên trình duyệt

Một điều thú vị về các mô hình ngôn ngữ lớn (LLM) dựa trên kiến trúc Transformer là trong giai đoạn tạo văn bản, chúng có khả năng lấy thông tin từ bất kỳ token nào trước đó. Tuy nhiên, mô hình cần phải có sự chọn lọc; nếu mọi token đều ảnh hưởng đến quá trình tạo văn bản như nhau, nó sẽ không mang lại hiệu quả cao. Quá trình này cần một cơ chế để quyết định mức độ ảnh hưởng của một token lên token tiếp theo.

Hóa ra, chúng ta hoàn toàn có thể trực quan hóa cơ chế này!

Bạn có thể nhấn hoặc di chuột qua bất kỳ token nào đã được tạo để xem các token trước đó đã "ảnh hưởng"* đến quá trình tạo văn bản như thế nào.

Đang tải... (Yêu cầu bật JavaScript)

* Từ "ảnh hưởng" có thể chưa hoàn toàn chính xác vì hình ảnh trực quan này đã được đơn giản hóa rất nhiều. Nó tính toán trọng số attention (attention weight), được điều chỉnh theo độ lớn của vector giá trị (value vector), tổng hợp trên tất cả các attention head và cộng dồn qua tất cả các lớp. Giá trị này sau đó được dùng để điều khiển độ mờ (opacity) của các token trước đó. Các giá trị lớn nhất luôn có độ mờ là 1 và các giá trị còn lại được nội suy.

Rất nhiều thông tin đã phải lược bỏ để giới hạn hình ảnh trực quan chỉ còn một giá trị số cho mỗi token trước đó. Chính vì vậy, khi bắt đầu triển khai, tôi đã nghĩ rằng nó có thể sẽ khó hiểu. Nhưng thực tế, nó lại tạo ra được một số mô hình (pattern) khá thú vị!

Ví dụ, trong prompt mặc định "Office Move Summary", bạn có thể di chuột qua các văn bản được sao chép nguyên văn như địa chỉ và ngày tháng. Bạn sẽ thấy dữ liệu gốc nổi bật lên khá rõ, bởi vì token được tạo ra đã lấy rất nhiều thông tin từ dữ liệu nguồn.

Điều này giải quyết một vấn đề mà trước đây tôi thấy khá khó hiểu về các LLM. Nếu chúng hoạt động bằng cách dự đoán các token tiếp theo dựa trên xác suất, tại sao chúng lại có khả năng sao chép và dán nội dung tốt đến vậy? Chẳng phải chúng sẽ sớm mắc lỗi do ngẫu nhiên sao?

Nhưng với cơ chế này, bạn có thể thấy rằng mô hình không dự đoán toàn bộ chuỗi từ một vài trạng thái nội bộ hạn chế. Vì có quyền truy cập vào tất cả các token trước đó, nó có thể quyết định chọn lọc thông tin từ token nào khi sao chép, nhờ đó xác suất xảy ra lỗi có thể rất thấp. Trong ví dụ "Debugging an Average Function", bạn có thể thấy mô hình khá nhỏ này (600 triệu tham số) có thể dễ dàng tái tạo toàn bộ một hàm JS ngoại trừ phần sửa đổi dự định. (Mặc dù thực tế nó không tự tìm ra vấn đề, nên vẫn cần một vài gợi ý.)

Một điểm thú vị khác là khi bạn di chuột qua từ "remain" trong câu "Existing access cards and phone numbers remain" ở prompt "Office Move Summary". Bạn có thể thấy nó lấy thông tin từ từ "work" trong câu "Existing employee access cards will work" và "stay the same" trong câu "company phone numbers will stay the same". Vì vậy, nó giống như đang kết hợp thông tin từ các từ trong cả hai cụm từ, điều mà tôi thấy khá hay.

Triển khai

Bản thân phần trực quan hóa là một ứng dụng React khá cơ bản sử dụng Transformers.js để tạo văn bản. Tuy nhiên, vì chúng ta cần trích xuất thêm dữ liệu từ mô hình để trực quan hóa, nên không thể sử dụng vòng lặp tạo văn bản thông thường. Tôi đã phải "vibe-code" (viết mã theo cảm tính/thử nghiệm) vòng lặp tạo văn bản trong ứng dụng để có thể theo dõi các giá trị cần trực quan hóa.

Mặc dù sử dụng một mô hình nhỏ hơn cho việc này, nó vẫn nặng hàng trăm megabyte, và việc chờ đợi nó tải xong trước khi hiển thị bất cứ thứ gì là không khả thi. Vì vậy, tôi đã tạo sẵn một loạt các prompt để có thể tải và xem ngay lập tức.

Một vấn đề hóc búa khác là một số thành phần trong hình ảnh trực quan thực tế không nhằm mục đích để đọc, nên chúng không được định nghĩa là các đầu ra (outputs). Tôi cho rằng nếu bạn triển khai việc này bằng các thư viện ML của Python, việc truy cập chúng vẫn sẽ dễ dàng. Nhưng Transformers.js sử dụng các tệp.onnx chứa toàn bộ đồ thị tính toán. Logic tải và tính toán mô hình được triển khai bằng wasm, nên theo những gì tôi biết, không có cách nào dễ dàng để truy cập bất cứ thứ gì ngoài các đầu ra đã được xác định trước.

Cuối cùng, tôi đã sử dụng một tập lệnh nhỏ để sửa đổi tệp onnx vừa đủ để làm lộ các giá trị nội bộ đó. Nhưng điều đó có nghĩa là tôi không thể chỉ sử dụng mô hình.onnx thông thường. Vì muốn có tính năng tạo văn bản trên trình duyệt, tôi phải tải một mô hình đã được chỉnh sửa (instrumented) riêng lên kho lưu trữ Hugging Face của mình và trỏ ứng dụng về đó.

Bạn có thể tìm thấy mã nguồn trong kho lưu trữ GitHub.