Có thể biến LLM thành máy tính không? — Cách chạy chương trình trực tiếp bên trong Transformer

• Để vượt qua giới hạn khi LLM có thể giải các bài toán ở mức Olympic Toán học nhưng vẫn không thể thực hiện chính xác cả phép cộng đơn giản/Sudoku nếu không có công cụ bên ngoài, một cách tiếp cận là xây dựng một máy tính thực sự bên trong Transformer
• Bằng cách chuyển đổi mã C tùy ý thành token, bản thân mô hình có thể tự tạo trực tiếp execution trace gồm hàng triệu bước, và có thể stream trên CPU với tốc độ hơn 30 nghìn token/giây
• Công nghệ cốt lõi là giới hạn chiều của attention head xuống 2D để chuyển sang truy xuất hình học dựa trên convex hull, thay thế phép quét thời gian tuyến tính bằng thời gian logarit
• Bằng cách triển khai trình thông dịch WebAssembly trong trọng số của Transformer, việc thực thi chương trình diễn ra minh bạch ngay trong vòng lặp giải mã của mô hình mà không cần gọi công cụ bên ngoài
• Về sau, cách làm này có thể mở rộng thành một hệ thống AI nơi chương trình được biên dịch trực tiếp thành trọng số, và biểu diễn đã học cùng thuật toán được biên dịch được hợp nhất trên một nền tảng tính toán duy nhất

---

Giới hạn tính toán của LLM

• Các LLM tối tân có thể giải bài toán ở mức huy chương vàng Olympic Toán học quốc tế hoặc thử thách các bài toán chưa được giải trong toán học và khoa học, nhưng vẫn thất bại ở các tác vụ tính toán thuần túy
• Ngay cả với phép cộng cơ bản cũng vẫn xảy ra lỗi, và ở các benchmark như Sudoku-Bench, tỷ lệ giải thành công nếu không có trợ giúp bên ngoài là rất thấp
• Hiện có hai cách lách vấn đề này
  • Sử dụng công cụ (Tool use): mô hình viết mã, trình thông dịch bên ngoài chạy mã đó rồi trả kết quả về
  • Điều phối agent: một vòng lặp bên ngoài lưu trạng thái trung gian, phân rã công việc và gọi lại mô hình nhiều lần
• Dù hữu ích, các cách tiếp cận này lại nhấn mạnh một giới hạn căn bản: năng lực tính toán thực sự nằm bên ngoài mô hình
• Có thể ví như con người chế tạo được máy bay không đồng nghĩa với việc con người có thể tự bay
• Mô hình có thể suy luận về phép tính hoặc điều phối công cụ, nhưng nếu không tự thực thi được thì chưa phải là một máy tính đúng nghĩa

Cách biến Transformer thành máy tính

• Tính phổ quát về mặt lý thuyết của kiến trúc Transformer trong việc mô phỏng máy Turing đã được chứng minh qua nhiều nghiên cứu, nhưng năng lực biểu diễn trên lý thuyết không đảm bảo hiệu quả thực thi trong thực tế
• Thay vì một mô hình tính toán thuần lý thuyết, nhóm nghiên cứu triển khai một máy tính RAM hiện đại bên trong Transformer, với mỗi lệnh được ánh xạ thành tối đa 5 token
• Tuy nhiên, vấn đề sâu hơn nằm ở chính quá trình giải mã
  • Giải mã autoregressive tiêu chuẩn ở mỗi bước đều phải tương tác với toàn bộ lịch sử ngày càng dài ra
  • KV cache tránh việc tính lại key/value trong quá khứ, nhưng chi phí attention vẫn tăng theo kích thước cache
• Để giải quyết giới hạn này, họ giới hạn chiều attention head xuống 2D và đưa vào một đường giải mã hiệu quả cho execution-style trace
  • Các phép truy xuất/cập nhật quan trọng có thể thực hiện trong thời gian logarit theo độ dài chuỗi
  • Nhờ đó có thể chạy các chương trình hàng triệu bước ngay bên trong Transformer

Ý nghĩa của tính toán — dùng công cụ vs thực thi trong mô hình

• Cách dùng công cụ truyền thống: mô hình in ra mã như python -c "print(3+5)" → trình thông dịch bên ngoài chạy → kết quả được bơm trở lại vào luồng token
• Cách của hệ thống này: triển khai trình thông dịch WebAssembly trong trọng số của Transformer
  • WebAssembly là một tập lệnh mức thấp cho thực thi nhanh và có tính quyết định, đồng thời là đích biên dịch phổ dụng cho C/C++
  • Khi tính 3 + 5, mô hình sẽ xuất ra các lệnh wasm (i32.const, i32.add) rồi chuyển sang chế độ giải mã nhanh để tự tạo execution trace từng bước
• Khác biệt cốt lõi
  • Dùng công cụ là mờ đục (opaque): mô hình bàn giao quyền kiểm soát rồi nhận lại câu trả lời như một hộp đen
  • Thực thi trong mô hình là minh bạch (transparent): mọi bước trung gian đều xuất hiện trong trace, và mô hình không rời khỏi vòng lặp giải mã của chính nó

Demo Sudoku — bài kiểm tra áp lực cho tính toán dài hạn chính xác

• Các cách tiếp cận mạng nơ-ron đã huấn luyện có thể cho kết quả tốt với Sudoku dễ, nhưng thất bại hoàn toàn ở bài khó
• Thường có lời giải thích rằng mô hình autoregressive vốn không phù hợp với bài toán thỏa mãn ràng buộc, nhưng hệ thống này vẫn hoàn toàn autoregressive mà vẫn đạt độ chính xác 100%
  • Nút thắt thực sự không nằm ở bản thân mô hình autoregressive, mà ở chỗ Sudoku khó đòi hỏi execution trace rất dài và attention tiêu chuẩn khiến việc tạo ngữ cảnh dài trở nên không thực tế
• Một bộ giải Sudoku hoàn chỉnh đã được biên dịch được chạy bên trong Transformer, thực hiện thuật toán thực sự từng bước chứ không phải heuristic đã học
• Giải chính xác Sudoku “khó nhất thế giới” của Arto Inkala trong chưa đầy 3 phút
• Nếu bộ giải đã biên dịch là chính xác, thì việc thực thi của Transformer cũng chính xác — cung cấp một bảo đảm phổ quát

Cách mã hóa phép tính — trace chỉ có thể nối thêm

• Có thể hình dung một Transformer autoregressive như một cỗ máy sống bên trong chính lịch sử của nó
  • Máy tính truyền thống cập nhật bộ nhớ có thể chỉnh sửa, nhưng Transformer không có khái niệm đó
  • Nó gồm một prompt cố định (đầu vào/chương trình) và một trace chỉ liên tục dài thêm ra (các token được tạo)
• Phép so sánh với cuốn sổ chỉ có thể ghi nối thêm
  • Những dòng đầu là đầu vào (prompt), sau đó mỗi dòng ghi lại bước tiếp theo của phép tính
  • Không thể sửa các dòng trước, và ở mỗi bước chỉ có thể tham chiếu đến một số ít vị trí trước đó
• Nhiều thuật toán có thể biểu diễn dưới dạng “một trace chỉ nối thêm, trong đó mỗi bước chỉ tham chiếu đến một số ít vị trí trước đó cố định”
• Trong hệ thống này, các token do mô hình tạo ra biểu diễn trạng thái đang tiến hóa của máy ảo như instruction pointer, thao tác bộ nhớ/ngăn xếp, số học, luồng điều khiển, đầu ra
• Attention head hoạt động như một mảng 1 chiều dùng chung, nơi mỗi token ghi giá trị vào một chỉ số và đọc giá trị từ chỉ số khác — cung cấp primitive ghi rồi đọc
• Attention cũng có thể tính cumulative sums, cho phép theo dõi instruction pointer, độ sâu ngăn xếp... dưới dạng tổng tích lũy của các delta tăng thêm

Mở khóa cốt lõi: attention nhanh theo cấp số mũ

• Máy tính thực sự cập nhật một trạng thái nhỏ với gần như lượng công việc hằng số cho mỗi lệnh, nhưng Transformer tiêu chuẩn ở bước giải mã thứ t phải tương tác với tiền tố độ dài t, khiến tổng chi phí tăng bậc hai
• HullKVCache được đưa vào để giải quyết sự bùng nổ bậc hai này
  • Trong benchmark thực tế, HullKVCache đạt 31.037 token/giây (6.747 dòng/giây), còn KVCache tiêu chuẩn đạt 272 token/giây (59 dòng/giây), chênh lệch khoảng 114 lần
  • Thay vì thời gian Θ(t) cho mỗi bước, việc truy xuất attention được thực hiện trong O(log t)

• Đường đi nhanh của attention 2D

  • Thay vì tăng tốc Transformer nói chung hay đưa vào kiến trúc mới, cách tiếp cận tập trung vào một tham số hóa dễ xử lý là giới hạn chiều head xuống 2D trong Transformer vanilla
  • Với d_model = 36, n_heads = 18, mỗi head có đúng 2 chiều, dùng 7 lớp, nn.MultiheadAttention tiêu chuẩn và mạng feedforward có cổng
    • Được xây dựng bằng PyTorch vanilla thuần túy, không cần kernel attention tùy biến hay mask thưa
  • Số lớp, số head và kích thước embedding vẫn có thể tăng tùy ý, nên toàn bộ mô hình không nhất thiết phải nhỏ
    • Dùng nhiều head hơn với n_heads = d_model / 2

• Góc nhìn hình học của attention 2D

  • Cơ chế attention: token quá khứ cung cấp vector key 2D kⱼ và value vⱼ, bước hiện tại tạo query q để trả về giá trị của key có tích vô hướng lớn nhất (hardmax attention)
  • Điều này tương đương chính xác với bài toán kinh điển trong hình học tính toán là “truy vấn điểm đỡ (supporting point query)”: cho hướng q, tìm điểm trên convex hull xa nhất theo hướng đó
  • Bằng cách duy trì cấu trúc dữ liệu convex hull ngay khi sinh token, có thể truy xuất trong thời gian logarit trên toàn bộ các điểm quá khứ
  • Với thao tác bộ nhớ/ngăn xếp, cần truy xuất “giá trị được lưu gần nhất tại chỉ số i”, và đó là lý do cần key 2D
    • Lưu chỉ số j dưới dạng key 2D kⱼ = (2j, -j²) và truy vấn bằng hướng q = (i, 1), khi đó hạng tử bậc hai -j² đóng vai trò phạt để chỉ phép khớp chính xác mới thắng
  • Attention 2D là đủ cho tính đầy đủ Turing, và như blog này cho thấy, nó có thể biểu diễn cả một máy tính RAM hoàn chỉnh

Kế hoạch tiếp theo

• Cơ chế attention phong phú hơn

  • Hiện tại triển khai dùng hardmax attention, nhưng đây không phải ràng buộc căn bản
  • Có thể xấp xỉ bằng k-sparse softmax attention: truy xuất top-k key từ các convex hull lồng nhau rồi chỉ chạy softmax trên chúng, với chi phí O(k + log n)
  • Đường nhanh hình học này không chỉ giới hạn ở cấu trúc bộ thực thi, mà về nguyên tắc có thể tăng tốc thời gian giải mã cho mọi Transformer có head 2D
  • Cũng có thể mở rộng tự nhiên sang head 3D thông qua convex hull 3D, nhưng ở chiều cao hơn hiệu quả suy giảm rất mạnh

• Huấn luyện mô hình lớn với head 2D

  • Tham số hóa head 2D về bản chất không hề nhỏ — vẫn có thể dùng nhiều head và nhiều lớp để duy trì tổng số tham số tương đương Transformer tiêu chuẩn
  • Có nhiều kiểu ứng dụng
    • Một đường nhanh chuyên dụng kết hợp với mô hình chậm hơn nhưng đa dụng hơn
    • Kiến trúc lai nhanh/chậm trong cùng một hệ thống
    • Speculative decoding, nơi mô hình 2D đề xuất token nhanh và mô hình attention thông thường xác minh lại
  • Vì execution trace là một phần của forward pass nên toàn bộ quá trình có thể vi phân (differentiable) — khác về bản chất so với công cụ bên ngoài, và là một nền tảng tính toán có thể học được để tích hợp trực tiếp vào mô hình lớn hơn

• Biên dịch chương trình thành trọng số

  • Hiện tại mô hình học một trình thông dịch được mã hóa trong trọng số, nhưng nếu mở rộng cỗ máy biên dịch sinh trọng số thì có thể biên dịch trực tiếp chương trình bất kỳ thành trọng số mà không cần chuỗi token
  • Bản thân trọng số sẽ trở thành đích triển khai (deployment target) của phần mềm, thay vì mô hình học cách hành xử như phần mềm thì logic chương trình đã biên dịch sẽ được nhúng thành một phần mạch nội bộ
  • Nếu có thể biên dịch logic vào trọng số, gradient descent sẽ không còn là cách duy nhất để sửa đổi mô hình — sẽ có thêm một con đường khác để chèn trực tiếp cấu trúc, thuật toán và bảo đảm vào mạng
  • Cuối cùng, điều này có thể phát triển thành các hệ thống không chỉ học từ trải nghiệm mà còn tự sửa đổi hoặc mở rộng trọng số của chính mình để tự viết lại cỗ máy nội bộ

• Làm cho hệ thống AI phát triển như phần mềm

  • Cũng như hệ sinh thái phần mềm hiện đại tiến hóa bằng cách tích lũy module, trừu tượng hóa và các thành phần tái sử dụng, một quá trình tương tự có thể diễn ra bên trong hệ thống AI khi năng lực tính toán mới được bổ sung dần dần
  • Chức năng mới có thể được nối vào các mạch hiện có mà không cần huấn luyện lại toàn bộ hệ thống
  • Các hệ thống AI tương lai sẽ không chỉ sử dụng phần mềm mà còn bao hàm phần mềm bên trong, hợp nhất biểu diễn đã học và thuật toán được biên dịch trên cùng một nền tảng tính toán

[GN⁺]

Ý kiến trên Hacker News

• Đây là một hướng tiếp cận thú vị hơn nhiều so với việc chỉ thực hiện tính toán
  Mô hình có thể thực hiện chuyển đổi attention động tỷ lệ với log của số token
  Nhờ vậy, nó có thể theo dõi các thanh ghi và ngăn xếp được biểu diễn bằng văn bản để mô phỏng việc thực thi chương trình
  Nếu LLM có thể chuyển sang một “chế độ tập trung” và tạo token cực nhanh, nó có thể tăng tốc giai đoạn suy luận vốn dùng để khám phá và sắp xếp vô số giả thuyết
  Bài báo đề xuất rằng điều này có thể được dùng cho kiến trúc hybrid kết hợp đường nhanh và đường chậm, hoặc làm mô hình speculative execution

• Lúc đầu tôi nghĩ “tại sao phải làm việc này?”, nhưng giờ tôi nhìn nó từ góc độ bootstrap huấn luyện
  Ví dụ, có thể nhúng một expert system đạt độ chính xác 80% vào mô hình, rồi dùng kết quả đó làm dữ liệu huấn luyện để nâng độ chính xác lên
  Càng giảm được chi phí huấn luyện cho nhiều tác vụ khác nhau thì rào cản gia nhập trong cuộc cạnh tranh AI càng thấp

  • Nhưng cách tiếp cận này không phù hợp cho huấn luyện
    Khác với softmax attention, average-hard attention không khả vi theo key và query
    Ngay cả khi hiệu chỉnh bằng straight-through estimator thì tốc độ backprop cũng không được cải thiện
  • Dù huấn luyện sẽ khó, vẫn có thể tham khảo bài báo này như một nghiên cứu liên quan thú vị

• Não người không mạnh về năng lực tính toán
  Nhân số 10 chữ số cũng mất nhiều thời gian. Những phép tính như vậy được cổng logic xử lý hiệu quả hơn nhiều
  Vì thế tôi tự hỏi liệu có tốt hơn không nếu LLM gọi một mô-đun bên ngoài kiểu ALU thay vì tự tính trực tiếp

  • Nhưng nếu nối chương trình nhúng với các trọng số khác của mô hình, ta có thể tích hợp nhiều thuật toán cố định đa dạng chứ không chỉ tính toán đơn thuần
  • Ngoài ra, cách tiếp cận này có thể là cơ chế để dạy cho mô hình trực giác hình học và khả năng suy luận không gian
  • Nếu còn chưa thử thì sẽ không biết được tiềm năng. Tích hợp tính toán xác định vào trong mạng nơ-ron cũng có thể giảm overhead gọi công cụ

• Văn phong có vẻ như do AI viết, nhưng thông điệp cốt lõi lại không rõ ràng
  Không rõ vì sao việc để mô hình chạy chương trình bên trong thay vì dùng hệ thống bên ngoài lại tốt hơn, và lợi ích nằm ở tốc độ, backprop hay benchmark

  • Tuy vậy, thay vì chỉ chê “trông như AI viết”, nên bàn về chính nội dung
    Một số người tin rằng symbolic logic là thiết yếu, nhưng nỗ lực kiểu này có thể đơn giản được xem là một thử nghiệm thú vị
  • Thực tế, nếu đọc bài báo thì execution trace là một phần của forward pass nên có thể khả vi, và gradient có thể được truyền qua chính phép tính đó
    Đây là khác biệt căn bản so với việc gọi công cụ bên ngoài
    Ngoài ra, nó có chi phí giải mã ở độ phức tạp O(k + log n), và nếu giải được các bài như Sudoku với độ chính xác 100% thì đã đủ ý nghĩa
    Cách này cũng có thể được kết hợp theo kiểu MoE, hoặc thử nghiệm nhiều trình thông dịch khác nhau như nhúng VM dựa trên WASM
  • Loại bỏ việc gọi công cụ bên ngoài còn giúp tăng tính bảo mật. Không còn rủi ro công cụ ngoài bị xâm phạm
  • Điểm cốt lõi là mô hình có thể viết và sửa mã trong lúc chạy. Đây là một kiểu thực thi động tương tự “aha moment” của con người
  • Sự lặp lại trong văn phong có thể chỉ là cách giải thích cho độc giả phổ thông. Con người cũng hay mắc lỗi như vậy

• Ý tưởng tích hợp công cụ vào đường tính toán nội bộ của mô hình rất thú vị ở góc độ khả năng diễn giải
  Thật đáng ngạc nhiên khi chỉ với Transformer đơn giản mà vẫn đạt được mức hiệu quả như vậy

• Có tiềm năng, nhưng ở trạng thái hiện tại thì tính thực dụng còn thấp
  Các trọng số hay công cụ “compiler” chưa được công bố nên rất khó thí nghiệm
  Dù vậy, ý tưởng nhúng các primitive tính toán được định nghĩa sẵn vào LLM vẫn có thể hữu ích

  • Nếu muốn hardcode các chương trình nhỏ vào trọng số Transformer, có thể tham khảo ALTA
  • Tôi tò mò không rõ cụm “neurosymbolic garbage” ở đây có nghĩa là gì

• Đây là câu then chốt:
  “Vì execution trace là một phần của forward pass, gradient có thể được truyền qua chính phép tính đó”
  Nói cách khác, khác với việc gọi công cụ bên ngoài, đây trở thành một nền tảng tính toán có thể huấn luyện được

  • Nhưng trên thực tế, đây không phải một cấu trúc khả vi hoàn toàn, và ngay trong bài báo cũng chỉ đề xuất các phương pháp xấp xỉ
    Ngoài ra, dữ liệu huấn luyện và thiết kế hàm loss cũng chưa rõ ràng
    Tuy nhiên, vì gọi công cụ làm vỡ hiệu quả batching, nên nếu cho đi qua một subnet tính toán nội bộ thì có thể cải thiện hiệu quả rất lớn ở quy mô lớn
    Dù vậy subnet đó chưa chắc phải là Transformer, có lẽ một layer không học tối ưu cho GPU là đủ

• Bài báo đang che giấu điểm cốt lõi
  Nếu giới hạn chiều của attention head xuống 2, có thể truy xuất và cập nhật token với độ phức tạp thời gian logarit
  Nhưng không rõ vì sao chiến lược này chỉ được áp dụng cho “code token”
  Việc chọn WASM làm đích cũng gây nghi ngờ về mặt hiệu quả

  • Với hai chiều, RoPE và hard-max attention, có thể hiện thực địa chỉ hóa tương đối chỉ bằng một head
  • Tuy nhiên, bài báo thiếu công thức và chi tiết huấn luyện, đồng thời dùng thuật ngữ phi chuẩn
    Ví dụ, gọi self-attention là “lookup table” là không chính xác
    Trong ví dụ mã, họ dùng d_model = 36, n_heads = 18 để tạo 2D per head, nhưng mọi thứ vẫn còn mơ hồ

• Không có giải thích cụ thể về việc họ đã compile Sudoku solver vào trọng số Transformer như thế nào
  Không rõ họ đã trực tiếp compile từ code sang weight, hay mô hình học được nó qua huấn luyện

  • Theo cách tôi hiểu, có vẻ họ đã nhúng một máy ảo đơn giản vào trọng số, rồi compile một runtime WASM lên trên đó để chạy solver
  • Thực ra trong bài báo có ghi rõ rằng họ đã huấn luyện một trình thông dịch WASM

• Thú vị đấy, nhưng vẫn còn câu hỏi “vì sao nhất định phải làm như vậy?”
  Não người cũng có thể mô phỏng máy Turing nhưng rất chậm. Vì vậy con người dùng công cụ bên ngoài
  Với mô hình cũng vậy, có lẽ dùng công cụ bên ngoài sẽ hiệu quả hơn

  • Bài báo đề xuất nhúng WebAssembly để giảm overhead của lệnh gọi Python, nhưng điều này khá giống tranh luận process vs thread từ thập niên 90
    Thậm chí cũng có thể nhúng một ngôn ngữ như Elixir để chạy mã ngắn hơn
  • Lập luận mang tính triết học rằng “một hệ thống không thể tính toán thì không thể hiểu tính toán” cũng khá thú vị
    Ý tưởng là mô hình có thể sửa mã khi đang chạy hoặc có khả năng debug
  • Nhưng bài báo không khám phá các khả năng đó mà chỉ dừng ở mức engine thực thi
  • Nếu đã định compile chương trình vào trọng số, thì tối ưu hóa việc gọi công cụ có lẽ còn hợp lý hơn
  • Ngoài ra cũng có thể dùng phép so sánh rằng nếu con người có thể nhúng máy tính cầm tay vào não thì sẽ nhanh và hiệu quả hơn nhiều