Kiến trúc Titans của Google: Giải thích các khái niệm then chốt

Năm 2017, Google đã làm rung chuyển cộng đồng AI với bài báo đột phá "Attention is All You Need." Bài báo này giới thiệu kiến trúc Transformer, vốn đã cách mạng hóa mô hình hóa chuỗi bằng cách sử dụng cơ chế tự chú ý (self-attention).

Gần đây, Google công bố một bài báo mới có tựa "Titans: Learning to Memorize at Test Time," trong đó họ đề xuất một kiến trúc bộ nhớ mới nhằm nâng cao cách hệ thống AI lưu trữ và khai thác thông tin.

Lấy cảm hứng từ sự tinh vi của trí nhớ con người, nghiên cứu này hướng đến việc khắc phục các hạn chế của hệ thống bộ nhớ hiện có trong các mô hình AI, đặc biệt là cửa sổ ngữ cảnh bị giới hạn của Transformers.

Kiến trúc Titans giới thiệu một cách tiếp cận tinh vi để tích hợp các loại bộ nhớ khác nhau—ngắn hạn, dài hạn và bền vững—vào các mô hình AI, cho phép chúng học và ghi nhớ một cách linh hoạt ngay trong thời gian kiểm thử.

Cơ chế "ngạc nhiên" là trung tâm của cách tiếp cận này, dựa trên xu hướng của não người ghi nhớ các sự kiện bất ngờ. Kiến trúc này cho phép các mô hình quản lý bộ nhớ hiệu quả bằng cách kết hợp những bất ngờ thoáng qua và trong quá khứ, bảo đảm việc ghi nhớ và truy xuất dữ liệu hiệu quả.

Trong bài viết này, tôi sẽ đơn giản hóa các khái niệm và cơ chế được giới thiệu trong bài báo Titans và khám phá tác động tiềm năng mà kiến trúc mới này có thể mang lại cho tương lai của AI.

Vấn đề với kiến trúc Transformer hiện tại

Transformers là kiến trúc được dùng trong hầu hết các LLM hiện nay. Chúng nổi tiếng với khả năng xử lý các chuỗi dữ liệu, như các từ trong câu. Chúng hoạt động rất hiệu quả nhờ một cơ chế gọi là "attention", giúp chúng tập trung vào những phần dữ liệu liên quan. Tuy nhiên, chúng đối mặt với một thách thức chủ chốt: kích thước cửa sổ ngữ cảnh.

Cửa sổ ngữ cảnh về cơ bản là lượng thông tin mà một Transformer có thể xem xét cùng lúc. Hãy nghĩ về nó như một cửa sổ chạy dọc theo một câu, nắm bắt một đoạn từ để hiểu ý nghĩa của chúng trong ngữ cảnh.

Lý tưởng nhất, chúng ta muốn cửa sổ này càng lớn càng tốt để xử lý những đoạn văn bản dài cùng lúc, nhưng có một trở ngại. Khi cố gắng mở rộng cửa sổ ngữ cảnh, lượng tính toán cần thiết tăng lên đáng kể do độ phức tạp thời gian và bộ nhớ bậc hai của kiến trúc transformer.

Điều này có nghĩa là nếu chúng ta tăng gấp đôi kích thước cửa sổ ngữ cảnh, khối lượng tính toán không chỉ tăng gấp đôi—mà trở nên nặng gấp bốn lần! Nguyên do là mỗi từ trong cửa sổ cần được so sánh với mọi từ khác, khiến số phép tính tăng vọt.

Các mô hình SOTA (tối tân) ngày nay có cửa sổ ngữ cảnh 2 triệu token. Tuy nhiên, chi phí xử lý chúng khá cao, và một số ứng dụng, như hiểu video, còn đòi hỏi cửa sổ ngữ cảnh lớn hơn. Có một cửa sổ ngữ cảnh dài hơn mà không làm bùng nổ chi phí tính toán sẽ cho phép LLM xử lý chuỗi dài hiệu quả hơn, cải thiện hiệu suất ở các tác vụ đòi hỏi hiểu những khối dữ liệu lớn cùng lúc. Đây là lý do vì sao khắc phục chi phí bậc hai là trọng tâm trong việc phát triển các mô hình này.

Kiến trúc Titans là gì?

Titans là một kiến trúc mới tăng cường khả năng lưu trữ và truy xuất thông tin của mô hình AI một cách hiệu quả, đặc biệt khi xử lý dữ liệu lớn và phức tạp, bằng cách lấy cảm hứng từ cơ chế trí nhớ con người.

Cảm hứng từ trí nhớ con người

Trí nhớ đóng vai trò then chốt trong việc học và hành vi của con người. Nó giúp chúng ta lưu giữ thông tin từ những trải nghiệm trước đó để ra quyết định, thích ứng với tình huống mới và giải quyết vấn đề. Trí nhớ của con người có thể được phân loại rộng thành ba loại:

• Trí nhớ ngắn hạn: Dùng để xử lý thông tin hiện tại. Nó giống như một không gian làm việc tinh thần, nơi ta tạm thời lưu và quản lý dữ liệu cần ngay lúc này, chẳng hạn như nhớ một số điện thoại vừa đủ lâu để bấm gọi.
• Trí nhớ dài hạn: Lưu trữ thông tin trong thời gian dài. Bao gồm kiến thức, kỹ năng và ký ức cá nhân, cho phép ta hồi tưởng trải nghiệm và học hỏi từ chúng ngay cả sau nhiều năm.
• Siêu trí nhớ (meta-memory): Hiểu biết về những gì ta biết. Nó giúp ta đánh giá mình có thể nhớ hay thực hiện một nhiệm vụ tốt đến đâu, từ đó ảnh hưởng đến quyết định và phán đoán.

Trong hành vi con người, các loại trí nhớ này phối hợp với nhau để cho phép chúng ta học từ môi trường, điều chỉnh chiến lược và cải thiện hành vi theo thời gian. Nếu không có chúng, ta sẽ bị giới hạn ở các phản xạ cơ bản và hành vi rập khuôn.

Cấu trúc bộ nhớ của Titans

Cũng như mô hình trí nhớ con người, cấu trúc bộ nhớ của Titan gồm ba phần:

• Trí nhớ ngắn hạn (lõi): Mô-đun này tập trung xử lý các nhiệm vụ hiện tại, cho phép AI xử lý dữ liệu tức thời với độ chính xác cao. Tương tự trí nhớ ngắn hạn của não, thành phần này giữ thông tin liên quan luôn sẵn sàng để xử lý nhanh nhưng không lưu giữ vĩnh viễn.
• Trí nhớ dài hạn: Đóng vai trò kho lưu trữ thông tin trong thời gian dài. Mô-đun này tương tự trí nhớ dài hạn ở người, lưu giữ trải nghiệm, kiến thức và ký ức không cần ngay lập tức nhưng có thể trở nên quan trọng sau này. Nó cho phép mô hình AI ghi nhớ và truy cập hiệu quả thông tin trong quá khứ, điều tối quan trọng cho các tác vụ đòi hỏi hiểu ngữ cảnh theo thời gian.
• Trí nhớ bền vững: Hoạt động như siêu trí nhớ của não, nhúng kiến thức liên quan đến nhiệm vụ vào bên trong mô hình. Loại trí nhớ này gồm các tham số không phụ thuộc vào đầu vào hiện tại nhưng rất quan trọng để hiểu và thực thi các nhiệm vụ cụ thể. Trí nhớ bền vững đảm bảo các khuôn mẫu và khung sườn đã học vẫn là một phần của mô hình, nâng cao khả năng áp dụng kiến thức cũ vào tình huống mới.

Bằng cách tích hợp ba loại trí nhớ này, kiến trúc Titans cho phép các mô hình AI quản lý thông tin một cách linh hoạt giống như bộ não chúng ta, hòa trộn xử lý tức thời với việc lưu giữ trải nghiệm và kiến thức quá khứ.

Cơ chế "ngạc nhiên"

Cơ chế "ngạc nhiên" trong Titans mô phỏng một khía cạnh nền tảng của trí nhớ con người: xu hướng ghi nhớ rõ ràng hơn các sự kiện bất ngờ. Cách tiếp cận này giúp mô hình AI quyết định thông tin nào cần lưu trữ bằng cách xác định dữ liệu khác biệt so với những gì mô hình kỳ vọng.

Hãy tưởng tượng bạn đang đi bộ trên phố và mọi thứ có vẻ bình thường, quen thuộc. Bạn có thể sẽ không nhớ chuyến đi này sau vài tuần. Tuy nhiên, nếu có điều gì đó bất ngờ xảy ra, như một cuộc diễu hành rực rỡ đi ngang qua hoặc nhận được một cuộc gọi với tin vui, những yếu tố bất ngờ này sẽ thu hút sự chú ý, khiến chúng đáng nhớ hơn.

Tương tự, Titans áp dụng ý tưởng về sự ngạc nhiên để ưu tiên thông tin cần ghi nhớ. Khi mô hình gặp dữ liệu khác biệt đáng kể so với dự đoán (giống như cuộc diễu hành rực rỡ kia), nó đánh dấu thông tin đó là gây ngạc nhiên. Hệ thống bộ nhớ trong Titans sau đó ưu tiên lưu trữ các sự kiện bất ngờ này, bảo đảm dữ liệu quan trọng nhưng bất thường không bị bỏ sót.

Cơ chế này cho phép Titans quản lý tài nguyên bộ nhớ hiệu quả bằng cách tập trung vào thông tin quan trọng, ngoại lệ thay vì cố gắng lưu mọi chi tiết. Giống như con người dễ nhớ điều bất ngờ hơn, Titans áp dụng chiến lược tự nhiên này để xử lý và lưu giữ thông tin quan trọng hiệu quả, giúp chúng thực hiện tốt hơn các nhiệm vụ phức tạp liên quan đến việc sàng lọc lượng dữ liệu khổng lồ.

Về mặt hình thức, khái niệm ngạc nhiên được định nghĩa là gradient của hàm mất mát theo đầu vào. Điều này về cơ bản đo lường mức độ các tham số của mô hình cần thay đổi trước đầu vào mới để tối thiểu hóa mất mát.

Trực giác đằng sau việc dùng gradient làm thước đo sự ngạc nhiên là gradient lớn ngụ ý cần những thay đổi đáng kể trong mô hình để tiếp nhận đầu vào mới, cho thấy đầu vào khác biệt với những gì mô hình từng thấy hoặc là điều bất ngờ, do đó gây ngạc nhiên.

Cơ chế quên

Cơ chế quên trong kiến trúc Titans là một thành phần quan trọng được thiết kế để quản lý bộ nhớ thần kinh dài hạn một cách hiệu quả, bảo đảm mô hình giữ lại thông tin liên quan đồng thời loại bỏ những gì không còn cần thiết.

Cơ chế này đặc biệt quan trọng khi xử lý lượng dữ liệu khổng lồ, nơi bộ nhớ có thể bị quá tải bởi thông tin lỗi thời hoặc không liên quan. Trên thực tế, cơ chế quên hoạt động bằng cách điều chỉnh quy tắc cập nhật bộ nhớ để xóa có chọn lọc thông tin dựa trên mức độ ngạc nhiên và mức độ liên quan, được đánh giá động trong quá trình suy luận.

Cụ thể, một tham số cổng thích ứng xác định giữ lại bao nhiêu trạng thái bộ nhớ trước đó, cho phép mô hình dọn dẹp hoặc bảo toàn nội dung bộ nhớ một cách thích nghi. Cơ chế như vậy ngăn tràn bộ nhớ và thúc đẩy sử dụng không gian bộ nhớ hiệu quả, cuối cùng nâng cao khả năng tập trung vào các đầu vào quan trọng và mới mẻ mà không bị đè nặng bởi sự tích tụ dữ liệu không cần thiết.

Cách tiếp cận này phản chiếu sát sao quá trình nhận thức của con người, trong đó những ký ức không quan trọng hoặc lỗi thời sẽ phai nhạt theo thời gian, cho phép xử lý và truy xuất thông tin hiệu quả hơn.

Học trong thời gian suy luận

Trong máy học truyền thống, việc học diễn ra trong giai đoạn huấn luyện. Khi mô hình được huấn luyện xong, nó được triển khai để suy luận—tạo dự đoán dựa trên dữ liệu mới, chưa từng thấy—mà không sửa đổi thêm các tham số đã học.

Kiến trúc Titans mang đến một bước chuyển đổi đột phá trong mô hình này bằng cách cho phép mô hình học một cách linh hoạt trong quá trình suy luận, còn gọi là thời gian kiểm thử. Khả năng này cho phép Titans cập nhật bộ nhớ và tham số dựa trên mức độ ngạc nhiên và tầm quan trọng của dữ liệu đến trong khi đang xử lý nó.

Khả năng học trong thời gian suy luận mang lại một số lợi thế đáng kể:

• Thích ứng linh hoạt với dữ liệu mới: Các mô hình truyền thống bị giới hạn bởi tính tĩnh của việc học, dựa trên các tham số cố định từ dữ liệu huấn luyện lịch sử. Bằng cách học trong lúc suy luận, Titans có thể thích ứng với các mẫu dịch chuyển và đầu vào mới theo thời gian thực, tăng mức độ phù hợp và độ chính xác khi đối mặt với môi trường thay đổi hoặc dữ liệu bất ngờ.
• Cải thiện mức độ liên quan của dữ liệu và quản lý bộ nhớ: Bằng cách liên tục cập nhật bộ nhớ trong thời gian kiểm thử, Titans có thể ưu tiên học từ dữ liệu được xem là gây ngạc nhiên hoặc khác biệt so với kiến thức hiện có. Việc quản lý bộ nhớ linh hoạt này phản chiếu khả năng của não người tập trung vào và ghi nhớ thông tin mới hoặc quan trọng trong khi loại bỏ điều thường nhật, nhờ đó tránh quá tải thông tin.
• Tăng cường khả năng khái quát hóa: Vì Titans có thể tinh chỉnh hiểu biết dựa trên dữ liệu trực tiếp, nó có thể khái quát hóa tốt hơn với dữ liệu ngoài phân phối—những đầu vào không có trong tập huấn luyện gốc. Khả năng này quan trọng trong các ứng dụng thực tế nơi đầu vào khác biệt đáng kể so với dữ liệu huấn luyện.
• Tăng năng lực xử lý phụ thuộc dài hạn: Học trong thời gian suy luận giúp Titans nhận diện và lưu trữ các phụ thuộc dài hạn trong dữ liệu hiệu quả hơn. Điều này đặc biệt có lợi cho các tác vụ cần xử lý chuỗi dài, như hiểu văn bản hoặc video dài.
• Tính ứng dụng thực tế cao: Cơ chế học trong thời gian suy luận cho phép Titans liên tục phát triển, học các mẫu và kiến thức mới khi chúng xuất hiện. Tính năng này phù hợp chặt chẽ với nhu cầu của các ứng dụng thực tế, nơi mô hình phải luôn phù hợp trong môi trường dữ liệu năng động và phức tạp.

Chế độ bộ nhớ trong Titans

Kiến trúc Titans sử dụng ba cách riêng biệt để tích hợp bộ nhớ, mỗi cách được lấy cảm hứng từ các khía cạnh khác nhau của nhận thức con người. Hãy cùng phân tích từng chế độ bộ nhớ bằng những so sánh đơn giản.

Nguồn: Behrouz et. al, 2024

Bộ nhớ như ngữ cảnh (MAC)

Hãy tưởng tượng bạn có một trợ lý cá nhân trong cuộc họp. Trợ lý này ghi chép đầy đủ các cuộc trao đổi trước đây (trí nhớ dài hạn) và thì thầm vào tai bạn những thông tin liên quan khi cần (truy xuất).

Ở chế độ MAC, mô hình AI hoạt động tương tự. Nó lưu giữ lịch sử chi tiết của dữ liệu quá khứ và tích hợp với đầu vào hiện tại để đưa ra quyết định sáng suốt. Giống như trợ lý giúp hòa trộn kiến thức cũ với thảo luận hiện tại để nâng cao chất lượng quyết định, MAC dùng trí nhớ dài hạn để cung cấp ngữ cảnh cho thông tin đang được xử lý.

Bộ nhớ như cổng (MAG)

Hãy hình dung bạn có hai cố vấn trong tâm trí. Một người chỉ tập trung vào những gì đang diễn ra, cung cấp góc nhìn cho các nhiệm vụ tức thời (chú ý ngắn hạn). Người còn lại rút ra từ nhiều năm kinh nghiệm và hiểu biết tích lũy (trí nhớ dài hạn).

Chế độ MAG đóng vai trò như người gác cổng quyết định nên dành bao nhiêu sự chú ý cho mỗi cố vấn, điều chỉnh động mức độ tập trung dựa trên nhu cầu hiện tại so với kiến thức đã học. Điều này đảm bảo cách tiếp cận cân bằng, cho phép AI thích ứng dù nhiệm vụ đòi hỏi độ chính xác tức thời hay cần đến insight từ kinh nghiệm quá khứ.

Bộ nhớ như tầng (MAL)

Hãy coi chế độ này như nhiều lớp hiểu biết, nơi mỗi lớp xử lý thông tin theo cách riêng:

• Lớp đầu có thể đóng vai trò bộ lọc lịch sử, tập trung vào các khuôn mẫu và ký ức dài hạn.
• Lớp thứ hai xử lý ngữ cảnh tức thời, như những quan sát ngắn hạn.
• Khi thông tin đi qua từng lớp, nó được tinh lọc và biến đổi, bảo đảm mô hình hiểu thấu đáo cả bối cảnh rộng và chi tiết tinh trước khi ra quyết định.

Quy trình theo lớp này phản chiếu cách chúng ta thường tiêu hóa thông tin phức tạp theo từng giai đoạn, xây dựng và tinh chỉnh hiểu biết từng bước. Mỗi chế độ mang lại thế mạnh riêng, giúp Titans xử lý hiệu quả các nhiệm vụ đa dạng bằng cách mô phỏng các khía cạnh khác nhau của cách con người xử lý và lưu giữ thông tin.

Hệ quả thực tế và thành công thực nghiệm

Kết quả thực nghiệm cho thấy sự vượt trội của kiến trúc Titans so với Transformers truyền thống và các mô hình hiện đại trong nhiều tác vụ. Dưới đây là tóm lược những phát hiện chính:

• Mô hình hóa ngôn ngữ và suy luận thường thức: Titans, đặc biệt biến thể bộ nhớ như ngữ cảnh (MAC), thể hiện hiệu năng vững vàng trong các tác vụ mô hình hóa ngôn ngữ và suy luận thường thức.
• Nhiệm vụ "mò kim đáy bể": Ở các tác vụ truy xuất mẩu thông tin cụ thể từ chuỗi dài, Titans vượt trội rõ rệt.
• Xử lý ngữ cảnh lớn hơn: Titans có khả năng mở rộng tới kích thước cửa sổ ngữ cảnh lớn hơn 2 triệu token mà không đánh đổi độ chính xác, là một cải tiến đáng kể so với các mô hình tối tân hiện nay.
• Dự báo chuỗi thời gian: Mô-đun bộ nhớ thần kinh trong Titans vượt trội so với nhiều đường cơ sở, bao gồm mô hình dựa trên Transformer và mô hình tuyến tính, trong các tác vụ dự báo.
• Mô hình hóa DNA: Kiến trúc của Titans cũng cho thấy hiệu năng cạnh tranh trong các tác vụ mô hình hóa DNA, chứng minh tính linh hoạt vượt ra ngoài các tác vụ ngôn ngữ tự nhiên truyền thống.

Tổng thể, Titans tận dụng cách quản lý bộ nhớ mới, cho phép chúng xử lý hiệu quả các tác vụ ngữ cảnh dài và phức tạp, thiết lập chuẩn mực mới cho mô hình AI cả về khả năng mở rộng lẫn độ chính xác.

Kết luận

Kiến trúc Transformer hiện tại, dù mạnh mẽ, vẫn gặp khó khi xử lý ngữ cảnh cực dài do vấn đề mở rộng. Kiến trúc sáng tạo Titans, được lấy cảm hứng từ trí nhớ con người, khắc phục những hạn chế này bằng cách quản lý linh hoạt trí nhớ ngắn hạn, dài hạn và bền vững.

Tôi tin rằng khả năng học trong thời gian suy luận, thích ứng với dữ liệu mới và ưu tiên bộ nhớ dựa trên cơ chế "ngạc nhiên" của Titans có tiềm năng tạo ra các hệ thống AI xử lý thông tin giống con người hơn. Điều này có thể cải thiện hiệu suất của hệ thống AI trong môi trường phức tạp và năng động, và có thể là bước ngoặt cho các tác nhân AI như Operator của OpenAI.