Dự đoán chuỗi đơn giản bằng transformer tự thiết kế thủ công, không cần huấn luyện

• Đây là một thử nghiệm tạo ra một decoder-only transformer giống GPT-2, nhưng không huấn luyện mà tự thiết kế trọng số để nó dự đoán mẫu (aab)*
• Việc dự đoán "aabaabaabaab..." cần nhìn vào hai token ngay trước đó, nên phù hợp hơn để bộc lộ cách attention hoạt động so với mẫu luân phiên đơn giản
• Mô hình được cấu hình nhỏ với N_CTX=5, N_VOCAB=2, N_EMBED=8, và dùng token hóa a=0, b=1 cùng embedding one-hot
• Một attention head duy nhất phân bổ 0.5 attention cho mỗi trong hai token gần nhất, rồi dùng triệt tiêu trong mã hóa a=1, b=-1 để tính token tiếp theo
• Trong các ngữ cảnh không mơ hồ, mô hình đạt độ chính xác 100.0% (27/27), nhưng cần khoảng 4.000 FLOPs để dự đoán một lần với ngữ cảnh 5 token, kém hiệu quả hơn rất nhiều so với quy tắc viết tay trực tiếp

---

Mini GPT-2 với trọng số được gán tay, không qua huấn luyện

• Mục tiêu là hiểu trực quan từng thành phần của transformer và attention thực sự làm gì
• Mô hình không được huấn luyện và cũng không dùng trọng số pretrained; thay vào đó, từng trọng số được gán thủ công trong một buổi tối
• Cấu trúc là một decoder-only transformer tương tự GPT-2, và phần triển khai được đơn giản hóa từ bản cài đặt picoGPT của jaymody
  • bỏ layer norm
  • dùng một head duy nhất thay vì multi-head attention
  • bỏ lớp feed-forward mlp trong transformer block

Vì sao chọn chuỗi (aab)*

• Ban đầu mục tiêu là dự đoán chuỗi như "ababababab", nhưng vì transformer dự đoán shifted sequence nên bài toán này quá dễ
  • chỉ cần dự đoán b nếu là a, ngược lại dự đoán a, nên không cần dùng positional embedding
• Bài toán cuối cùng được chọn là dự đoán chuỗi "aabaabaabaab...", tức (aab)*
  • nếu hai token trước đó là ab hoặc ba thì token tiếp theo là a
  • nếu hai token trước đó là aa thì token tiếp theo là b
  • bb được xem là trường hợp ngoài phạm vi bài toán
• Token hóa cũng rất đơn giản vì chỉ xử lý hai ký hiệu
  • a là 0
  • b là 1

Kích thước mô hình và luồng tính toán

• Có ba tham số mô hình được chọn
  • N_CTX = 5: độ dài ngữ cảnh tối đa mô hình nhìn cùng lúc
  • N_VOCAB = 2: hai token a, b
  • N_EMBED = 8: kích thước embedding chứa token, vị trí và không gian tính toán
• Dù bài toán thực tế chỉ cần hai token trước đó, N_CTX=5 được chọn để bao gồm cả tình huống phải bỏ qua các token không liên quan
• Hàm gpt hoạt động theo thứ tự sau
  • cộng token embedding và positional embedding bằng wte[inputs] + wpe[range(len(inputs))]
  • đi qua một transformer block
  • cuối cùng tạo logits trong không gian từ vựng bằng x @ wte.T

Embedding chứa vị trí và token dưới dạng one-hot

• wpe là positional embedding, dùng 5 chiều embedding đầu tiên làm one-hot vị trí
  • position 0 là [1, 0, 0, 0, 0]
  • position 4 là [0, 0, 0, 0, 1]
• wte là token embedding, dùng 2 chiều tiếp theo làm one-hot token
  • token a là [1, 0] tại các chiều token tương ứng
  • token b là [0, 1]
• Vị trí embedding thứ 8 ban đầu chưa dùng đến, và được dùng làm scratch space bên trong transformer block
• Ví dụ, "aabaa" được biểu diễn thành một ma trận embedding 5 x 8, trong đó mỗi hàng chứa cả one-hot vị trí và one-hot token

Cách attention head chọn hai token gần nhất

• Transformer block gồm một attention head và một lớp tuyến tính c_proj đưa kết quả attention trở lại không gian embedding
• c_attn là lớp tuyến tính kích thước embed_size x (embed_size * 3), biến embedding đầu vào thành ma trận qkv, rồi tách ra thành q, k, v
• k tách positional embedding ra để biểu diễn thông tin vị trí mà mỗi token mang theo
• q biểu diễn phạm vi vị trí mà mỗi vị trí muốn tìm, và tạo ma trận điểm attention qua q @ k.T
• Sau softmax và causal mask, ma trận attention có các tính chất sau
  • hàng đầu tiên chỉ attention 100% vào token đầu tiên
  • các hàng sau đó chia attention 0.5 cho mỗi trong hai token gần nhất mà có thể truy cập
• Causal mask cộng thêm một giá trị rất nhỏ vào vị trí token tương lai, trong mã thực tế là giá trị như -1e10, để không thể nhìn token tương lai
  • mô hình làm tay này không được thiết kế để nhìn trộm tương lai, nhưng mask vẫn được giữ lại để gần với cấu trúc GPT-2 hơn
• Phần scaling bằng np.sqrt(q.shape[-1]) giúp cải thiện gradient trong huấn luyện thực tế, nhưng không ảnh hưởng đến mô hình thủ công này

Mã hóa v và phép triệt tiêu cộng để tạo dự đoán

• v biến one-hot token thành mã hóa a=1, b=-1
• Vì kết quả attention lấy trung bình 0.5 của hai token gần nhất, cách mã hóa này tính ra các quy tắc sau
  • a, b → 0.5 * 1 + 0.5 * (-1) = 0
  • b, a → 0.5 * (-1) + 0.5 * 1 = 0
  • a, a → 0.5 * 1 + 0.5 * 1 = 1
• Kết quả là ở vị trí thứ 7 của hàng sẽ xuất hiện các giá trị sau
  • 0 nếu cần dự đoán a
  • 1 nếu cần dự đoán b
• Với đầu vào "aabaa", hàng đầu tiên có thể sinh ra dự đoán b do thiếu thông tin, nhưng các dự đoán sau đó khớp với quy tắc (aab)*

Đưa giá trị dự đoán trở lại không gian từ vựng

• c_proj biến giá trị ở vị trí thứ 7 của kết quả attention trở lại dạng one-hot token
• Nó không đơn giản tạo [..., 1, 0, ...] hoặc [..., 0, 1, ...], mà tạo one-hot được scale bởi 1024
  • embedding[row, 5] = 1024 + (-1024) * prediction
  • embedding[row, 6] = 0 + 1024 * prediction
• Transformer block có residual connection, nên embedding gốc được cộng lại qua x = x + causal_self_attention(...)
• Vì residual signal còn sót lại không cần thiết, scale 1024 được dùng để lấn át tín hiệu này
• Cuối cùng tính x @ wte.T để tạo logits rồi áp dụng softmax
  • trong ngữ cảnh "aabaa", hàng dự đoán cuối cùng sẽ chỉ về b
  • trong lúc huấn luyện, dự đoán ở mọi hàng đều hữu ích, nhưng khi suy luận chỉ cần hàng cuối cùng

Kết quả sinh và độ chính xác

• Hàm complete đưa tối đa 5 token cuối vào mô hình, rồi chọn token tiếp theo bằng argmax trên hàng cuối của kết quả softmax
• Ví dụ kết quả sinh như sau
  • complete("a") → a :: baabaabaab
  • complete("ba") → ba :: abaabaabaa
  • complete("abaab") → abaab :: aabaabaaba
• Ngay cả với đầu vào ngoài phạm vi, đôi khi mô hình vẫn phục hồi về mẫu lặp
  • complete("ababa") → ababa :: abaabaabaa
  • complete("bbbbb") → bbbbb :: aabaabaaba
• Trên bài kiểm tra "aab" * 10, nếu chỉ đánh giá các ngữ cảnh không mơ hồ thì độ chính xác là 100.0% (27/27)

Khác biệt giữa 4.000 FLOPs và 8 lệnh

• Khi dùng toàn bộ ngữ cảnh 5 token, mô hình này cần khoảng 4.000 floating point operations để dự đoán một token
  • phần lớn nằm ở phép tính attention
  • có thể giảm bằng cách thu nhỏ context window, fused multiply-add, kv caching, v.v.
  • dù vậy, để dự đoán một token vẫn cần tới hàng trăm lệnh máy
• Cùng quy tắc (aab)*, một đoạn x64 assembly viết tay chỉ cần 8 lệnh để tính token tiếp theo
• Điều đó để lại câu hỏi liệu có thể huấn luyện một language model hiệu quả hơn mô hình hiện tại tới 1000 lần cho việc sinh ngôn ngữ tự nhiên hay không

[GN⁺]

Ý kiến trên Hacker News

• Một công trình liên quan là "Thinking Like Transformers"
  Bài này giới thiệu một ngôn ngữ lập trình nguyên thủy tên là RASP, gồm các phép toán có thể được mô hình hóa bằng các thành phần của Transformer, và cho thấy có thể viết các chương trình như histogram hay sắp xếp
  Cũng có các bài blog xuất sắc của Sasha Rush và Gail Weiss; các nghiên cứu tiếp theo đã cho thấy các chương trình kiểu RASP có thể được biên dịch thành trọng số mô hình thực tế mà không cần huấn luyện
  [1] https://arxiv.org/abs/2106.06981
  [2] https://srush.github.io/raspy/
  [3] https://arxiv.org/abs/2301.05062

  • Mình rất thích dòng RASP
    Nếu thấy lĩnh vực này thú vị, cũng đáng xem công trình HandCrafted Transformers của tôi, trong đó tự chọn trọng số cho mô hình Transformer để nó thực hiện phép cộng dài theo cách tương tự như con người học ở tiểu học
    [1] https://colab.research.google.com/github/newhouseb/handcraft...
  • Có vẻ ngôn ngữ hàm như Haskell sẽ rất phù hợp với những công việc kiểu này
    Hướng đi từ mạng nơ-ron sang mã cũng có lẽ rất thú vị ở khía cạnh khả năng giải thích

• Tôi tưởng mình hiểu khá rõ Transformer, nhưng chưa từng tự triển khai
  Một ngày nọ tôi thử tự triển khai thì nó không hoạt động hay học tốt như PyTorch Transformer chuẩn, và cuối cùng nhận ra nguyên nhân là mình đã bỏ qua dropout
  Tôi cho nó học phép cộng số và chưa bao giờ cho thấy cùng một cặp hai lần, nên nghĩ là không thể overfit, nhưng vai trò của dropout hóa ra lớn hơn tôi tưởng nhiều
  Tóm lại, cứ tự triển khai Transformer là tốt, và càng làm từ nền tảng càng tốt
  Tất cả những người đã làm đều học được điều gì đó ngoài dự đoán; mỗi người ngộ ra những phần khác nhau, từ song song hóa học theo đơn vị token cho đến backpropagation thực sự hoạt động như thế nào

  • Không biết có tài liệu tham khảo nào hữu ích để tiếp cận công việc này không

• Tài liệu của Karpathy cũng hay, nhưng video khiến tôi cuối cùng hiểu được Transformer là video này: https://youtu.be/kWLed8o5M2Y?si=SJT5_lCJ0hSR7Z_k

• Tôi đã nghĩ về một ý tương tự trong một thời gian
  Liệu có thể tạo ra một giao diện trực quan cho trọng số mô hình, để chuyên gia miền có thể tinh chỉnh thủ công nhằm tăng tốc huấn luyện không?
  Ví dụ trong mô hình thị giác, khi phát hiện cọc tiêu giao thông, ta tăng một nhóm trọng số tương ứng với "độ cam"
  Khi đó con người có thể tăng tốc thay vì yêu cầu thêm hàng nghìn hay hàng triệu ví dụ để hiệu chỉnh đúng "độ cam"
  Dĩ nhiên điểm khó là giao diện này phải ánh xạ được tới các nhóm trọng số mang các ý nghĩa khác nhau, và tôi tò mò liệu có lý do kỹ thuật nào khiến việc đó bất khả thi không

  • "Chuyên gia miền tinh chỉnh thủ công trọng số của mô hình" nghe giống với cách nhận dạng ảnh trước thời deep learning
    [1] https://www.youtube.com/watch?v=8SF_h3xF3cE&t=1358s
  • Lý do bạn đang tìm được gọi là The Bitter Lesson
    Nói ngắn gọn, cách để con người hỗ trợ AI hầu như luôn kém hiệu quả về chi phí hơn so với việc chạy bằng nhiều sức mạnh tính toán hơn
    Trong lúc con người hiệu chỉnh một lớp trọng số để phát hiện cọc tiêu giao thông màu cam, cụm GPU đã huấn luyện AI phát hiện cọc tiêu giao thông, đèn tín hiệu, cây cối, xe khác, và cả các cọc tiêu giao thông có sắc cam hơi khác
  • Số lượng lớp và trọng số không ở quy mô mà con người có thể cập nhật thủ công; mà dù có thể đi nữa thì việc quản lý hiệu ứng hạ nguồn khi thay đổi trọng số cũng quá khó
    Dù bạn điều chỉnh hình ảnh để nhìn màu cam tốt hơn, nếu không thể đồng thời giám sát độ chính xác của mọi màu khác, rất có thể bạn vô tình tạo ra vấn đề ở các màu khác
  • Lý do việc này bất khả thi hoặc rất khó về mặt kỹ thuật là vì trọng số thường rất khó diễn giải
    Không phải một cụm neuron cụ thể tương ứng với một khái niệm cụ thể; nhìn chung mọi thứ đều làm mỗi việc một chút
  • Cơ chế attention của Transformer có vẻ không dễ ánh xạ sang ngữ nghĩa mà con người có thể hiểu
    Có quá nhiều tham số liên quan

• Bài báo về Transformer quá kỹ thuật nên tôi luôn muốn hiểu dù chỉ ở mức nông, nhưng thấy khó
  Bài viết này thực sự giúp tôi hiểu cách nó hoạt động, và ít nhất các ví dụ thì rất rõ ràng
  Nhờ vậy tôi cũng nhớ lại được ma trận đã học hồi đại học

• Chẳng phải nó là một dạng máy trừu tượng, giống như máy Turing hay máy phân tích cú pháp biểu thức chính quy sao?

  • Nếu đơn giản hóa một chút, nó là một “cỗ máy” ánh xạ tập đầu vào thành tập xác suất của đầu ra tiếp theo.
    Trước hết định nghĩa danh sách token, ví dụ cho dễ hiểu là 24 ký tự.
    Cỗ máy này nhận một chuỗi token đầu vào, thực hiện các phép toán ma trận mang tính quyết định, rồi xuất ra danh sách xác suất của mọi token.
    “Huấn luyện” chỉ là quá trình thiết lập một phần các con số trong những ma trận dùng cho các phép toán đó.
    Đáng chú ý là trong mã cuối cùng chỉ có một câu lệnh if, mà câu lệnh đó cũng chỉ để đánh giá độ chính xác của kết quả.
    Toàn bộ “logic” đều đến từ kết quả của các phép toán ma trận.
  • Khá khó để diễn giải những thứ như vậy thành automata theo nghĩa mà ta thường nghĩ đến.
    Trong mạng nơ-ron, mọi thứ nhìn chung đều hơi mờ, và gần như không có những thứ như if/else, dù vẫn có các trường hợp “mask” giá trị thành 0 hoặc -∞ như trong ví dụ Transformer.
    Đầu ra hầu như luôn là một nhóm điểm số hoặc xác suất, nên nếu một mô hình phân biệt ảnh chó và mèo cho kết quả kiểu dog:0.95 cat:0.05, ta nói nó dự đoán là chó vì điểm của chó cao hơn.
    Cơ chế attention, phần cốt lõi của Transformer, dựa trên một dạng phép toán tra cứu mềm.
    Với một hệ thống không mờ, ta sẽ duyệt từng token trong chuỗi, kiểm tra xem nó có liên quan đến token hiện tại không, rồi nếu có thì thực hiện một hành động nào đó; nhưng trong Transformer, mức độ liên quan không phải là một quyết định nhị phân.
    Thay vào đó, nó tính các điểm liên quan liên tục giữa mọi cặp token trong chuỗi, rồi dùng các điểm đó để quyết định hành động tiếp theo.
    Tuy vậy, có những thứ không dễ tổng quát hóa trực tiếp từ một hệ thống dựa trên quyết định nhị phân.
    Chẳng hạn, các điểm liên quan đó được dùng làm trọng số để tính trung bình có trọng số trên các token từ vựng, từ đó thu được một “token trung bình” cho vị trí hiện tại.
    Có vẻ không có cách nào dễ dàng diễn giải việc này như một phần mở rộng của quy trình dựa trên logic rẽ nhánh.
  • Chẳng phải là AllSpark đã chạm vào một đống đại số tuyến tính sao?
  • Đúng vậy.
    Nên xem bài báo này, giải thích rằng Linear Transformers thực ra là Fast Weight Programmers: https://arxiv.org/abs/2102.11174
  • Mạng nơ-ron là máy Turing.
    Nếu thiết lập trọng số một cách cẩn thận, ta có thể khiến nó thực hiện bất kỳ phép tính nào.
    Chỉ là sẽ tốt hơn nếu có một compiler không dựa trên xấp xỉ.

• Tôi tò mò câu “có thể bạn sẽ muốn tự tạo một mô hình” định dùng vào việc gì ngoài một bài tập học tập để thỏa mãn trí tò mò.
  Tôi bắt đầu có cảm giác các mô hình machine learning phức tạp là thứ không thực tế đối với người ngồi nhà đọc blog.

  • Trong nanoGPT, nếu tiền huấn luyện mô hình bằng Shakespeare, chỉ sau 3 phút nó đã đạt mức trung thành với dữ liệu gốc ngang tầm Jabberwocky của Lewis Carroll.
    Nó tạo ra nhiều từ tiếng Anh cổ trông có vẻ hợp lý, học được nền tảng ngữ pháp tiếng Anh và hình thức của kịch bản, v.v.
    Tôi thấy khá bất ngờ khi nó đi được xa đến vậy trong thời gian ngắn như thế.
    Huấn luyện cục bộ nhiều mô hình đến mức trung thành kiểu Shakespeare-from-Wish.com có lẽ sẽ giúp đánh giá xem ta đã tìm được kiến trúc tốt hay chưa và đã đến lúc thử mở rộng hay chưa.
  • Mục đích đã được nêu ở đoạn đầu của bài viết.
    Tác giả muốn hiểu Transformer và attention rõ hơn; đã đọc The Illustrated Transformer nhưng vẫn chưa có trực giác về việc các mảnh khác nhau của attention thực sự làm gì.
    Đại ý là còn chưa phân biệt rõ q và k, chưa nói đến v.
  • Đây là một bài tập học tập rất hay.
    Nó không chỉ thỏa mãn trí tò mò mà còn giúp xây dựng và đào sâu sự hiểu biết.
  • Cũng có thể đơn giản là có người thật sự thích hack những dự án kiểu này thì sao? Hơi kỳ, nhưng có thể mà.

• Ước gì tiêu đề có thể thêm cụm như neural network.
  Nội dung này liên quan đến kiến trúc “Transformer” trong machine learning, chứ không phải một bó cuộn dây ghép hai mạch bằng điện từ.