Cải tạo mạng bộ nhớ ngắn hạn dài: XLSTM cho AI thế hệ tiếp theo

LSTM hay Mạng bộ nhớ ngắn hạn dài đã xuất hiện từ lâu. Chúng đã được áp dụng cho khá nhiều tác vụ liên quan đến trình tự, chẳng hạn như tạo và dịch văn bản hoặc thậm chí tạo chú thích cho hình ảnh.

Hạn chế của chúng là không thể song song hóa để sử dụng các GPU mạnh mẽ hiện đại. Hạn chế này đã mở đường cho sự xuất hiện của các máy biến áp tận dụng GPU để thực hiện song song hóa quá trình đào tạo và suy luận.

Nếu bây giờ chúng ta cố gắng cải tiến và song song hóa các LSTM, liệu chúng có thể trở thành công cụ để xây dựng LLM thế hệ tiếp theo không?

Đây là câu hỏi chính xác được trả lời bởi bài báo “ XLSM - Mạng bộ nhớ ngắn hạn dài mở rộng “, viết tắt của bộ nhớ ngắn hạn dài “ mở rộng “. Họ làm như vậy bằng cách đề xuất hai khối mới trong kiến trúc, đó là sLSTM và mLSTM.

Vì vậy, hãy đi sâu vào các khối sLSTM và mLSTM được đề xuất trong bài viết này và xem cách chúng ta có thể xếp chúng lại với nhau để phát triển kiến trúc XLSTM.

Giải thích trực quan

Nếu bạn là người như tôi và muốn XLSTM được giải thích một cách trực quan thì vui lòng xem video YouTube về bài viết này:

Bồi dưỡng LSTM

Một trong những mạng đầu tiên được thiết kế để xử lý dữ liệu tuần tự là Mạng thần kinh tái phát.

Mạng thần kinh tái phát

Nó sử dụng các kết nối lặp lại trong kiến trúc của nó với x là đầu vào và o là đầu ra. Nếu mở nó ra, chúng ta có thể hình dung nó như một chuỗi các thao tác xảy ra tại các mốc thời gian t-1, t và t+1 . Một nhược điểm lớn của RNN là vấn đề biến mất độ dốc, trong đó độ dốc bằng 0 khi chúng ta xếp quá nhiều khối lại với nhau.

LSTM, hay Mạng bộ nhớ ngắn hạn dài, đã được đề xuất để khắc phục độ dốc biến mất bằng cách đưa các trạng thái ô và cơ chế cổng vào mạng.

Một minh họa đơn giản về LSTM

Các trạng thái tế bào, c , là những ký ức dài hạn tồn tại qua nhiều dấu thời gian. Các trạng thái ẩn, h , là những ký ức ngắn hạn được truyền từ bước này sang bước khác. Và tất nhiên, chúng ta có đầu vào z từ chuỗi đầu vào.

Ba cổng có chức năng hình chữ S. Cổng quên sử dụng hàm sigmoid để quyết định thông tin nào cần quên trong bộ nhớ dài hạn. Cổng đầu vào cũng sử dụng hàm sigmoid để xử lý đầu vào và thêm nó vào đầu ra của cổng quên. Hoạt động bổ sung này có một thuật ngữ ưa thích được gọi là băng chuyền lỗi không đổi trong bài báo XLSTM và tài liệu học thuật. Hoạt động bổ sung này giải quyết vấn đề độ dốc biến mất được tìm thấy trong RNN. Sau đó, c_t đầu ra được xử lý bởi cổng đầu ra , thường là hàm tanh dẫn đến đầu ra trạng thái ẩn h_t được chuyển sang bước tiếp theo.

Với các thao tác này, chúng tôi đã mổ xẻ hai phương trình chính của LSTM là c_t và h_t .

Hạn chế 1 - Sửa đổi các quyết định lưu trữ

Một trong những nhược điểm chính của LSTM là không có khả năng sửa đổi các quyết định lưu trữ. Điều này có nghĩa là khi độ dài chuỗi tăng lên, mô hình sẽ có thể quyết định liệu nó có lưu giữ thông tin quá khứ trong bộ nhớ hay không.

Ví dụ: nếu chúng ta lấy câu này, “Tom đã đến cửa hàng. Anh ấy đã mua một ít đồ uống” và so sánh nó với “Tom đến cửa hàng để mua một số mặt hàng tạp hóa bao gồm cà rốt, hành tây, chuối, táo, cam, cà phê và bánh mì. Anh ấy cũng mua một ít đồ uống. Đối với mỗi từ mới, chẳng hạn như chuối hoặc táo, mô hình phải liên tục xem xét lại liệu nó có nên giữ từ “Tom” trong quá khứ trong bộ nhớ hay không. Đây là một thách thức lớn đối với LSTM và nó bắt nguồn từ chức năng sigmoid của cổng quên của nó.

Hàm sigmoid và hàm số mũ. Đường sigmoid phẳng dần về phía cuối, nhưng số mũ chỉ tiếp tục tăng.

Vì vậy, nếu chúng ta lấy một cổng quên, nó bao gồm một hàm sigmoid, có đường cong hình chữ S dẹt về phía cuối. Điều này cho thấy rằng khi chúng ta tiến tới các giá trị đầu vào cao hơn, việc quyết định cái gì nên quên và cái gì cần giữ trong bộ nhớ trở nên khá khó khăn. Nhưng nếu chúng ta sử dụng hàm mũ thay thế, thì trò chơi sẽ thay đổi và khi chúng ta chuyển sang giá trị đầu vào cao hơn, chúng ta sẽ có phạm vi rộng hơn cho đầu ra. Ngược lại, điều này chỉ ra rằng LSTM có thể hiệu quả hơn trong việc sửa đổi các quyết định lưu trữ.

Giải pháp 1 - sLSTM

Vì vậy, giải pháp được đề xuất trong bài báo này là các khối hệ thống. Nếu chúng ta quay lại phương trình LSTM cổ điển biểu thị trạng thái ô, như chúng ta đã thấy trước đây, thì đó là một hàm của cổng quên và cổng đầu vào.

Ngược lại, các cổng này bao gồm các hàm sigmoid. Vậy điều gì sẽ xảy ra nếu chúng ta thay thế các hàm sigmoid này bằng các hàm mũ? Các cổng mới f_t và i_t bây giờ trở thành exp(f_t) và exp(i_t), và đó gần như là sửa đổi chính để tạo khối sLSTM.

Không giống như hàm sigmoid, vốn ép các đầu vào ở một phạm vi cố định, các hàm mũ có xu hướng tăng giá trị khi đầu vào tăng và nó không chuẩn hóa đầu ra một cách tự nhiên để nằm giữa 0 và 1 như sigmoid chức năng.

Vì vậy, chúng ta cần giới thiệu một trạng thái chuẩn hóa mới, đây là chức năng của cổng quên và cổng nhập. Chúng ta có thể coi nó như một giá trị trung bình đang chạy của một giá trị chuẩn hóa.

Chúng tôi sử dụng các giá trị chuẩn hóa được tính toán để chuẩn hóa đầu ra hoặc trạng thái ẩn mới.

Trong khi quá trình chuẩn hóa xử lý các trạng thái ẩn, để kiểm soát hàm mũ từ việc làm nổ tung các cổng quên và cổng nhập, chúng ta cần giới thiệu một bộ ổn định . Nó xuất hiện dưới dạng hàm log để chống lại tác động của hàm mũ và mang lại sự ổn định. Vì vậy, trạng thái ổn định là giá trị lớn nhất của log của cổng quên và đầu ra của cổng đầu vào. Chúng tôi trừ các giá trị bộ ổn định này khỏi đầu vào và quên các cổng để ổn định chúng.

Hạn chế 2 - Bộ nhớ và song song

Hạn chế thứ hai của LSTM là thiếu tính song song. LSTM được thiết kế để xử lý dữ liệu tuần tự, có nghĩa là nó cần đầu ra xử lý đầu vào trước đó trong chuỗi để xử lý đầu vào hiện tại trong chuỗi. Hạn chế đặc biệt này ngăn cản sự song song hóa và là thủ phạm dẫn đến buổi bình minh của kỷ nguyên Transformers.

Giải pháp được đề xuất trong bài viết này là các khối mLSTM mới. Vì vậy, chúng ta hãy nhìn vào chúng tiếp theo.

Giải pháp — mLSTM

Khối xây dựng tiếp theo của XLSTM là khối mLSTM, trong đó m là viết tắt của bộ nhớ. Chúng ta hãy quay lại phương trình LSTM cổ điển một lần nữa để xem nhược điểm của nó là gì. Chúng ta có thể thấy rằng trạng thái ô c_t là vô hướng. Điều này có nghĩa là chúng tôi chỉ xử lý 1 số tại một thời điểm khi chúng tôi có sự sang trọng của GPU hiện đại với bộ nhớ ít nhất 12 Gigs.

Khối mLSTM đưa ra các ma trận thay cho các đại lượng vô hướng cho các trạng thái ô. Quay trở lại phương trình LSTM cổ điển của chúng ta, Điều gì sẽ xảy ra nếu chúng ta thay thế c_t bằng ma trận C *_t* để trạng thái ô bây giờ trở thành C viết hoa *_t* để biểu thị ma trận và trạng thái ô có thể được truy xuất không chỉ bằng cổng i_t mà còn bằng cách lưu trữ các cặp khóa-giá trị là vectơ. Các giá trị của nó có thể được truy xuất bằng các truy vấn là vectơ có cùng thứ nguyên.

Để làm cho nó quen thuộc với thuật ngữ của máy biến áp, họ đã giới thiệu khóa và giá trị ở đây để tạo thành ma trận này.

XLSTM

Với thông tin đó về sLSTM và mLSTM, hãy đi sâu vào kiến trúc chi tiết của XLSTM.

sLSTM

Cái nhìn chi tiết về khối sLSTM

Khi nói đến sLSTM, chúng tôi sử dụng các phép chiếu hậu kỳ. Vì vậy, đầu vào trước tiên được chuyển qua các lớp tích chập nhân quả với chức năng kích hoạt chuyển động. Đầu ra từ các lớp này sau đó được đưa qua một lớp tuyến tính chéo khối với bốn khối chéo hoặc “đầu”. Đầu ra từ những khối này sau đó được đưa qua khối sLSTM có bốn đầu. Cuối cùng, đầu ra được chiếu lên bằng cách sử dụng lớp MLP có kiểm soát với kích hoạt GeLU và được chiếu xuống bằng cách sử dụng chức năng MLP có kiểm soát.

mLSTM

Chuyển sang chi tiết của khối mLSTM, chúng tôi sử dụng các phép chiếu chuẩn bị trước. Có nghĩa là đầu vào được chiếu lên trước tiên với hệ số chiếu là 2. Một trong các đầu ra chiếu đi đến mLSTM và đầu ra còn lại đi đến cổng đầu ra. Đầu vào của khối mLSTM đi qua tích chập nhân quả và sau đó đi qua các ma trận chiếu chéo khối có kích thước khối 4, tạo ra truy vấn, khóa và giá trị mà khối mLSTM dễ dàng sử dụng.

Kiến trúc XLSTM

Cuối cùng, chúng ta có thể xếp chồng hai loại khối này để tạo thành kiến trúc LSTM mở rộng. Vì vậy, các khối màu xám đậm là các khối mLSTM và các khối màu xám nhạt là các khối sLSTM.

Về ưu điểm, bài báo đề cập rằng mạng XLSTM có độ phức tạp tính toán tuyến tính và độ phức tạp bộ nhớ không đổi liên quan đến độ dài chuỗi.

Đánh giá

Các tác giả đã đào tạo về bộ dữ liệu SlimPajama để so sánh nó với các phương pháp dựa trên máy biến áp khác như LLAMA và các phương pháp dựa trên không gian trạng thái như MAMBA. Họ đã sử dụng ký hiệu xLSTM a:b, trong đó a là số khối mLSTM và b là số khối sLSTM trong ngăn xếp.

Về độ chính xác, họ báo cáo độ chính xác tương đối bằng cách chia tỷ lệ độ chính xác trong khoảng từ 0 đến 1 trong đó 0 là ngẫu nhiên và 1 là hoàn hảo.

Các đánh giá chỉ ra rằng XLSTM hoạt động tốt hơn trong các nhiệm vụ như Parity, trong khi Llama và Mamba hoạt động kém hơn.

Từ các kết quả, một điều được đặc biệt quan tâm là nhiệm vụ tính chẵn lẻ, trong đó các máy biến áp hoặc các mô hình không gian trạng thái có xu hướng gặp khó khăn khi không trộn lẫn bộ nhớ hoặc theo dõi trạng thái. Chúng ta có thể thấy rằng trong loại nhiệm vụ này, xLSTM đạt độ chính xác bằng 1 khi chúng ta sử dụng cả hai khối sLSTM và mLSTM cùng nhau.

Họ cũng đã thực hiện một số nghiên cứu cắt bỏ để chứng minh tính bền vững của XLSTM. Họ rất dễ hiểu từ bài báo. Hơn nữa, bài viết này nói nhiều hơn về những điểm mới lạ trong kiến trúc của XLSTM, vì vậy tôi sẽ không đi sâu vào kết quả thử nghiệm ở đây.

La lên

Nếu bạn thích bài viết này, tại sao không theo dõi tôi Twitter nơi tôi chia sẻ thông tin cập nhật nghiên cứu từ các phòng thí nghiệm AI hàng đầu mỗi ngày trong tuần?

Ngoài ra hãy đăng ký theo dõi của tôi Kênh Youtube nơi tôi giải thích các khái niệm và bài viết về AI một cách trực quan.

Phần kết luận

Hy vọng bài viết này đã đơn giản hóa và giúp bạn hiểu rõ hơn về kiến trúc XLSTM, lý do tại sao chúng ta cần chúng và làm thế nào chúng có thể vượt qua máy biến áp trong tương lai gần.

Hãy chờ xem họ có gì trong cửa hàng. Tôi sẽ gặp bạn trong lần tiếp theo của tôi…