Mixture-of-Depths: kỹ thuật phân bổ động tài nguyên tính toán trong Transformer

 (arxiv.org)
2 điểm bởi GN⁺ 2024-04-08 | 1 bình luận | Chia sẻ qua WhatsApp
  • Mixture-of-Depths (MoD) của Google DeepMind là phương pháp cho phép mô hình ngôn ngữ Transformer không dùng cùng một lượng FLOPs cho mọi token, mà ở mỗi layer chỉ cho những token cần thiết tham gia vào tính toán attention và MLP
  • Router theo từng layer tạo trọng số vô hướng cho mỗi token; chỉ các token nằm trong sức chứa top-k được định trước mới đi qua block, còn các token còn lại đi vòng qua bằng residual connection
  • Bằng cách cố định trước k, đồ thị tính toán và kích thước tensor vẫn được giữ tĩnh, trong khi phân bổ tính toán theo token vẫn thay đổi động theo ngữ cảnh
  • Trong thử nghiệm, theo tiêu chí isoFLOP, MoD đạt loss thấp hơn Transformer cơ bản trong cùng thời gian huấn luyện, hoặc giảm FLOPs trên mỗi forward pass ở cùng mức hiệu năng, giúp tăng tốc các step huấn luyện và sampling
  • Thiết lập tốt nhất là áp dụng block sức chứa 12,5% cho các block cách nhau một block (every other block); trong auto-regressive sampling, phương pháp này giảm vấn đề phi nhân quả của top-k bằng routing dựa trên predictor

Lãng phí tính toán mà MoD nhắm tới

  • Mô hình ngôn ngữ Transformer thông thường dùng cùng một lượng tính toán cho mọi token trong forward pass
  • MoD dựa trên tiền đề rằng không phải mọi token và mọi chuỗi đều cần cùng thời gian và nỗ lực để dự đoán; phương pháp này giảm tổng ngân sách tính toán trong khi phân bổ tính toán cho các token cần thiết
  • Conditional computation hiện có có thể tạo đồ thị tính toán động, nên có thể không phù hợp với đồ thị tính toán tĩnh mà phần cứng hiện nay ưa chuộng
  • Khi người dùng xác định tổng ngân sách tính toán trước khi huấn luyện, mô hình sẽ học cách sử dụng ngân sách đó ở đâu theo từng token và từng layer
    • Có thể dự đoán trước lượng tính toán, mức dùng bộ nhớ và FLOPs trên mỗi forward pass
    • Các token được routing thay đổi theo ngữ cảnh đầu vào

Cấu trúc routing và cách giảm tính toán

  • MoD dùng router giống MoE Transformer, nhưng thay vì chọn một trong nhiều expert, nó gửi token vào một trong hai đường
    • Tính toán block tiêu chuẩn gồm attention và MLP
    • Residual connection chuyển nguyên giá trị đi tiếp
  • Residual connection có chi phí tính toán thấp, và đầu ra của block được quyết định bởi giá trị đầu vào
  • Nếu đặt capacity của block nhỏ hơn tổng số token T, FLOPs trên mỗi forward pass sẽ giảm so với Transformer cơ bản
    • Ví dụ, nếu giảm capacity xuống T/2, phép nhân ma trận query-key của attention sẽ là (T/2)² thay vì T², khiến FLOPs còn khoảng 25%
  • Nếu đặt capacity bằng 0 thì rất nhanh nhưng hầu hết tham số Transformer không được dùng nên hiệu năng giảm; nếu đặt capacity bằng T thì giống Transformer cơ bản
  • Mục tiêu là tìm điểm nằm giữa hai cực này, nơi có thể tăng tốc độ step trong khi đạt hiệu năng bằng hoặc tốt hơn Transformer cơ bản

Routing top-k theo expert-choice

  • Ở mỗi layer l, router tạo trọng số vô hướng cho từng embedding token
  • Theo capacity C do người dùng đặt, chỉ các token thuộc top-k trong chuỗi mới tham gia tính toán attention và MLP
  • Các token không được chọn sẽ bỏ qua tính toán block và được chuyển sang layer tiếp theo qua residual connection
  • Cách này dựa trên expert-choice routing
    • token-choice routing có thể phát sinh vấn đề load balancing vì token tự chọn đường đi
    • expert-choice routing cho phép mỗi đường chọn các token top-k, nên capacity được lấp đầy chính xác
  • Trong MoD, đường đi về cơ bản chỉ gồm hai loại: “tính toán block” và “đi vòng”, nên có thể dùng một phép top-k duy nhất để chia token thành hai tập
  • Trọng số router được nhân với đầu ra block, nằm trên gradient path và chịu ảnh hưởng của gradient descent trong quá trình huấn luyện language modeling

Vấn đề phi nhân quả khi sampling

  • Phép top-k của expert-choice routing là phi nhân quả (non-causal)
    • Việc một token có nằm trong top-k hay không chỉ được quyết định khi biết cả trọng số router của các token phía sau
    • Trong auto-regressive sampling, không thể dùng thông tin token tương lai
  • Để tránh điều này, nhóm đã thử nghiệm hai phương pháp
    • Auxiliary loss binary cross-entropy để sigmoid của đầu ra router dự đoán việc có được chọn vào top-k hay không
    • Một MLP predictor phụ nhỏ nhận cùng đầu vào với router nhưng áp dụng stop-gradient
  • Phương pháp đầu tiên ảnh hưởng khoảng 0,2~0,3% đến objective language modeling chính
  • Phương pháp thứ hai dựa trên predictor không ảnh hưởng đến objective language modeling và cũng không ảnh hưởng đáng kể đến step speed
  • Tác vụ dự đoán phụ trong thực nghiệm nhanh chóng đạt độ chính xác 99%

Kết quả thử nghiệm và quan sát

  • Các thử nghiệm ban đầu được tiến hành bằng cách tìm siêu tham số MoD trong tổng ngân sách 6e18 FLOPs
  • Thiết lập MoD tối ưu là áp dụng routing ở mỗi block cách nhau một block, với top-k đặt là 256
    • Với độ dài chuỗi 2048, chỉ 256 token, tức 12,5%, đi qua attention và MLP
    • 1792 token còn lại, tức 87,5%, đi vòng qua block
  • Biến thể MoD 220M tham số đạt hiệu năng nhỉnh hơn baseline tối ưu isoFLOP một chút, đồng thời step speed trong huấn luyện nhanh hơn hơn 60%; tổng thời gian huấn luyện wall-clock trên cùng phần cứng gần như tương đương
  • Trong phân tích isoFLOP ở 6e18, 2e19 và 1e20 FLOPs, các mô hình FLOP-optimal của MoD cũng có xu hướng có nhiều tham số hơn baseline
  • Tồn tại các biến thể MoD đạt cùng hiệu năng hoặc loss thấp hơn trong khi có FLOPs trên mỗi forward pass ít hơn
  • Cải thiện step speed đến từ hai yếu tố
    • Một số token đi vòng qua block, nên FLOPs per parameter giảm ở cùng kích thước mô hình
    • Có thể chọn biến thể MoD nhỏ hơn nhưng vẫn đạt hiệu năng như baseline
  • MoD dùng stochastic routing cho hiệu năng thấp hơn đáng kể so với Transformer cơ bản và MoD thông thường, cho thấy routing được học là quan trọng
  • Trong đánh giá auto-regressive, nhóm dùng held-out data gồm 256.000 chuỗi và 500M token; khi chuyển sang routing dựa trên predictor, mức suy giảm hiệu năng là nhỏ
  • MoD cũng có thể kết hợp với MoE để tạo thành Mixture-of-Depths-and-Experts (MoDE)
    • staged MoDE quyết định gửi token vào block hay đi vòng trước attention
    • integrated MoDE tích hợp routing MoD bằng cách thêm expert “no-op” vào giữa các MLP expert hiện có
  • Phân tích routing cho thấy một số token thường xuyên đi qua nhiều block, còn các token khác thì đi vòng bất cứ khi nào có thể
  • Phân tích sơ bộ cho thấy các token thường xuyên đi qua block có tương quan với trường hợp entropy của dự đoán đầu ra cao hơn, điều này liên quan đến khả năng đó là các token khó dự đoán hơn

Bài viết liên quan

1 bình luận

 
GN⁺ 2024-04-08
Ý kiến trên Hacker News

  • Định tuyến phức tạp hơn chắc chắn sẽ trở nên phổ biến hơn
    Đặc biệt, tôi nghĩ một ngày nào đó sẽ đi theo hướng định tuyến đệ quy, tức đưa tập chuyên gia đi qua lại lần nữa. Trong tương lai, có lẽ 'chuỗi suy nghĩ (chain-of-thought)' sẽ diễn ra đệ quy bên trong mô hình

    • Có thể gọi đối tượng giả định như vậy là Recursive Neural Networks
    • Điều được mô tả ở đây trông hơi giống dòng nghiên cứu Universal Transformers. Cách này đưa embedding đầu vào đi qua cùng một khối transformer nhiều lần, rồi một mô-đun riêng sẽ đánh giá xem embedding đã đủ “chín” hay chưa để lấy ra
      Gần với ý tưởng “chuyên gia” hơn là bài báo Sparse Universal Transformers năm ngoái, kết hợp Universal Transformer với hỗn hợp chuyên gia thưa, trong đó cơ chế gating quyết định dùng khối transformer nào theo thứ tự nào
      Đây không phải lĩnh vực chuyên môn của tôi, nhưng theo tôi hiểu thì khá khó huấn luyện cho đúng, và để đạt kết quả tương tự transformer thông thường thì khi suy luận cần tổng lượng tính toán lớn hơn. Dù vậy đây vẫn là hướng thú vị, và việc số bước tính toán trên mỗi token có giới hạn trên là một trong những nhược điểm lớn của kiến trúc transformer cổ điển
    • Tôi nghĩ lý do việc này vẫn chưa làm được là vì tại thời điểm huấn luyện không có cách quyết định cần đệ quy bao nhiêu lần
      Nếu chọn số lần ngẫu nhiên hoặc thử nhiều độ sâu đệ quy, đầu ra sẽ bị “nhòe”. Nói cách khác, không còn biết đầu ra của lớp nào nên cung cấp thông tin quan trọng cho kết quả cuối cùng, hay nên cung cấp đầu vào tốt nhất cho vòng đệ quy tiếp theo
    • Attention về cơ bản là định tuyến, và các kiểu định tuyến khác như thế này có thể cung cấp cho mô hình những lựa chọn thô hơn, qua đó có khả năng giúp việc học dễ hơn
    • Xu hướng rõ ràng là tăng định tuyến động, nhưng tôi cho rằng MoE/MoD/MoDE thiên về việc chứa thêm dữ kiện trong trọng số bằng cách giảm chồng lấn, hơn là giúp suy luận sâu hơn
      Suy luận sâu hơn có lẽ sẽ đến từ tính động ở cấp token hơn là tính động ở cấp lớp. Ví dụ có bài báo Quiet-STaR gần đây, trong đó mô hình xuất ra các token lập luận rồi loại bỏ chúng: https://arxiv.org/abs/2403.09629

  • Có thể là bài báo quan trọng nhất năm 2024
    Ý tưởng muốn có mô hình không cần dùng cùng một lượng tính toán cho mọi token đã có từ lâu, nhưng đây là lần đầu tôi thấy một cơ chế thuyết phục cho việc đó

    Equipped with these new methods, we can sample autoregressively by choosing to route tokens to or around a block based on the router’s output, which does not depend on any information from future tokens. We provide empirical evidence that this is a relatively easy auxiliary task that quickly achieves 99% accuracy.
    Chẳng phải điều này hơi đáng kinh ngạc sao?

    • Sparse Universal Transformer đã có từ lâu hơn và đã thực hiện dừng sớm dựa trên định tuyến rồi
    • Đến mức quan trọng nhất ư? Việc không phải mọi token đều cần toàn bộ cửa sổ ngữ cảnh lẽ ra phải là một tối ưu hóa hiển nhiên

  • Phiên bản nhập môn đơn giản:
    Hãy tưởng tượng có một trợ lý thông minh có thể hiểu và xử lý lời nói. Thông thường, trợ lý này dành cùng mức chú ý cho mọi từ, bất kể mỗi từ quan trọng thế nào đối với ý nghĩa tổng thể
    Bây giờ giả sử ta đã tìm ra cách dạy trợ lý dùng “tài nguyên não bộ” thông minh hơn. Thay vì dành cùng mức chú ý cho mọi từ, nó tập trung hơn vào những từ quan trọng nhất để hiểu ý nghĩa, và điều chỉnh sự tập trung đó ngay tại chỗ tùy theo ngữ cảnh
    Để trợ lý không bị quá tải, ta cũng đặt giới hạn cho tổng “tài nguyên não bộ” có thể dùng tại một thời điểm nhất định. Giống như cấp cho nó một ngân sách và nói rằng “mỗi lần chỉ được dùng tài nguyên cho một số từ nhất định”. Khi đó trợ lý phải quyết định từ nào là quan trọng nhất
    Dù có giới hạn này, trợ lý vẫn linh hoạt điều chỉnh việc sử dụng tài nguyên. Nó dùng nhiều hơn cho một số từ, ít hơn cho những từ khác, và đặt ưu tiên phù hợp với tình huống
    Khi được huấn luyện như vậy, trợ lý hoạt động hiệu quả mà vẫn chú ý một cách thông minh. Nó hiểu tốt như một trợ lý dành chú ý như nhau cho mọi từ, nhưng dùng ít tài nguyên não bộ tổng thể hơn, nên phản hồi và xử lý thông tin mới nhanh hơn nhiều

    • Tôi hiểu đây là ELI5, nhưng cách được mô tả chẳng phải đã là việc attention đang làm sao? Nó tập trung cụ thể vào những từ quan trọng theo ngữ cảnh trong chuỗi trước đó mà

  • Tôi đã tóm tắt một chút ở đây dựa trên những gì mình hiểu:
    https://lifeinthesingularity.com/p/googles-breakthroughs-in-...

    • Bài viết hay. Gợi nhớ phong cách New Scientist. Có cảm giác “vừa đi vừa giải thích nhưng vẫn giữ ngắn gọn”, rất tốt để nắm bắt lĩnh vực này

  • Rất giống hỗn hợp chuyên gia. Chỉ khác là thay vì định tuyến token tới nhiều chuyên gia, nó “đưa vào một chuyên gia duy nhất có thể được bỏ qua động”

    • Nếu trộn hai thứ lại thì có vẻ khá hay. Có thể giảm thêm lượng tính toán MoE trong khi vẫn giữ hiệu năng

  • “This is more computationally efficient than performing a full content-based lookup across an entire memory buffer for each step in the future, and could be one step towards drastically increasing the context-length available for making a prediction.”
    Đây có phải là cách giúp có cửa sổ ngữ cảnh 10 triệu token không? Hay là nói về cửa sổ ngữ cảnh dài hơn trong tương lai?

  • Kết luận tôi rút ra sau khi cố hiểu và triển khai vài thuật toán RASP là một số hàm nhất định cần một số lớp transformer nhất định để hoạt động
    Theo logic này, có thể thấy các hàm mà transformer học được có thể phân tán trên nhiều head. Việc lặp lại những hàm như vậy có thể rất có giá trị trong việc hiểu và giải quyết vấn đề, nhưng cách suy luận hiện nay không thể chạy lặp lại một nhóm head liên tiếp. Bài báo này thật sự có vẻ là một hướng đầy hứa hẹn
    [1] https://arxiv.org/pdf/2106.06981.pdf
    [2] https://www.youtube.com/watch?v=t5LjgczaS80

  • Nhược điểm duy nhất của việc LLM phát triển quá nhanh có lẽ là các bài báo ra nhanh hơn tốc độ mà bất kỳ ai, ít nhất là bên ngoài Google, có thể học và kiểm thử các cải tiến
    Khi tôi bắt đầu với deep learning, ReLU và dropout đang là xu hướng, và có thể sửa một hai dòng code trên GPU 1080 tiêu dùng rồi kiểm tra trong vài giờ xem có cải thiện hay không. Giờ có lẽ phải đợi vài tuần cho đến khi những nơi như Mistral thử nghiệm

    • Chào mừng đến với tầng lớp nghèo GPU
      Tôi đang tập trung vào các cách tiếp cận lượng tử hóa và kiểm thử trên những GPU thế hệ trước đã cũ

  • hu-po review chuyên sâu các bài báo AI bằng livestream
    Rất khuyến nghị, và ở đây có bàn về bài báo mixture-of-depths đang được thảo luận: https://www.youtube.com/watch?v=Teru_qIdB8Y