Published on

Kiến trúc Titan của Google phá vỡ nút thắt bộ nhớ Transformer

Tác giả
  • avatar
    Tên
    Ajax
    Twitter

Giới thiệu Titan: Kiến trúc mới từ Google

Thế giới công nghệ đang xôn xao về Titan, một kiến trúc mới nổi lên từ Google. Nó được thiết kế để thách thức những hạn chế của các mô hình Transformer, đặc biệt là cách chúng xử lý bộ nhớ. Kiến trúc mới này đang thu hút sự chú ý đáng kể như một người kế nhiệm tiềm năng của Transformer, đặc biệt khi nó được phát triển bởi một nhóm trong Google.

Thách thức về bộ nhớ trong các mô hình hiện tại

Các mô hình truyền thống như LSTM và Transformer, mặc dù rất đổi mới, phải đối mặt với những thách thức trong việc mô phỏng bộ nhớ giống con người. Những thách thức này bao gồm:

  • Dung lượng hạn chế: Dữ liệu thường được nén vào một trạng thái ẩn có kích thước cố định, hạn chế lượng thông tin có thể được giữ lại.
  • Chi phí tính toán: Mặc dù có khả năng nắm bắt các phụ thuộc tầm xa, chi phí tính toán tăng theo cấp số nhân với độ dài chuỗi, khiến nó không hiệu quả đối với các chuỗi rất dài.
  • Quá phụ thuộc vào dữ liệu huấn luyện: Việc chỉ ghi nhớ dữ liệu huấn luyện không phải lúc nào cũng giúp ích cho ứng dụng thực tế, nơi dữ liệu thử nghiệm có thể nằm ngoài phân phối huấn luyện.

Cách tiếp cận của Titan: Mô-đun bộ nhớ lấy cảm hứng từ thần kinh

Nhóm Titan đã có một cách tiếp cận khác, tìm cách mã hóa thông tin vào các tham số của mạng nơ-ron. Họ đã phát triển một siêu mô hình trực tuyến được thiết kế để học cách ghi nhớ và quên dữ liệu cụ thể trong quá trình thử nghiệm. Mô hình này được lấy cảm hứng từ các nguyên tắc tâm lý thần kinh, kết hợp các yếu tố chính sau:

  • Bất ngờ như một yếu tố kích hoạt: Các sự kiện bất ngờ dễ nhớ hơn. "Bất ngờ" được đo bằng độ dốc của đầu vào đối với mô-đun bộ nhớ. Độ dốc càng lớn, đầu vào càng bất ngờ.
  • Cơ chế động lượng và quên: Cơ chế động lượng tích lũy những bất ngờ ngắn hạn vào bộ nhớ dài hạn, trong khi cơ chế quên sẽ xóa những ký ức cũ, ngăn chặn tình trạng tràn bộ nhớ.
  • Bộ nhớ dựa trên Multi-Layer Perceptron (MLP): Mô-đun bộ nhớ bao gồm nhiều lớp MLP, cho phép nó lưu trữ các trừu tượng sâu về dữ liệu, khiến nó mạnh hơn bộ nhớ dựa trên ma trận truyền thống.

Cách tiếp cận siêu học trực tuyến này giúp mô hình tập trung vào việc học cách thích ứng với dữ liệu mới, thay vì chỉ ghi nhớ dữ liệu huấn luyện. Mô-đun này cũng được thiết kế để tính toán song song, nâng cao hiệu quả của nó.

Tích hợp mô-đun bộ nhớ vào các kiến trúc Deep Learning

Nhóm nghiên cứu Titans đã đề xuất ba biến thể để kết hợp mô-đun bộ nhớ của họ vào các kiến trúc deep learning:

  1. MAC (Memory as Context): Phương pháp này kết hợp bộ nhớ dài hạn và bộ nhớ liên tục (mã hóa kiến thức nhiệm vụ) làm ngữ cảnh được đưa vào cơ chế chú ý.
  2. MAG (Memory as Gate): Cách tiếp cận này sử dụng hợp nhất có cổng của mô-đun bộ nhớ với cơ chế chú ý cửa sổ trượt trên hai nhánh.
  3. MAL (Memory as Layer): Ở đây, mô-đun bộ nhớ được triển khai như một lớp độc lập, nén thông tin lịch sử trước khi đưa nó vào cơ chế chú ý.

Nhóm nghiên cứu nhận thấy rằng mỗi biến thể đều có những điểm mạnh và điểm yếu riêng.

Hiệu suất và ưu điểm của Titans

Titans đã chứng minh hiệu suất vượt trội trong nhiều tác vụ, bao gồm mô hình hóa ngôn ngữ, lập luận thông thường và dự đoán chuỗi thời gian. Nó đã vượt qua các mô hình hiện đại như Transformer và Mamba. Đáng chú ý, mô-đun bộ nhớ dài hạn (LMM) một mình đã vượt trội hơn các mô hình cơ sở trong một số tác vụ, thể hiện khả năng học tập độc lập của nó mà không cần bộ nhớ ngắn hạn (chú ý).

Trong một thử nghiệm "kim trong đống cỏ khô" được thiết kế để tìm các manh mối chi tiết trong các văn bản dài, Titans duy trì độ chính xác khoảng 90% ngay cả khi độ dài chuỗi tăng từ 2k lên 16k. Nhóm nghiên cứu chỉ ra rằng các thử nghiệm tiêu chuẩn không thể hiện đầy đủ lợi thế của Titans trong việc xử lý các văn bản dài. Titans cũng vượt trội hơn các mô hình như GPT4, Mamba và thậm chí Llama3.1 với RAG trong một nhiệm vụ yêu cầu suy luận từ các sự kiện trải rộng trên các tài liệu cực kỳ dài.

Titans cũng đã thể hiện hiệu suất ấn tượng trong các lĩnh vực cụ thể như dự đoán chuỗi thời gian và mô hình hóa chuỗi DNA.

Đội ngũ đứng sau Titans

Nghiên cứu được thực hiện bởi một nhóm từ nhóm thuật toán và tối ưu hóa của Google Research NYC, hiện không thuộc Google DeepMind.

  • Ali Behrouz, một thực tập sinh tại Đại học Cornell, là tác giả đầu tiên của bài báo.
  • Zhong Peilin, cựu sinh viên Đại học Thanh Hoa và là tiến sĩ tốt nghiệp Đại học Columbia, là nhà khoa học nghiên cứu tại Google từ năm 2021. Ông nổi tiếng với việc xuất bản một bài báo là tác giả đầu tiên tại STOC 2016 khi còn là sinh viên đại học.
  • Vahab Mirrokni, Google Fellow và VP, dẫn dắt nhóm.

Nhóm đã phát triển Titans bằng Pytorch và Jax và có kế hoạch sớm phát hành mã để đào tạo và đánh giá.