StreamingLLM - Triển khai mô hình ngôn ngữ streaming hiệu quả với attention sink
(github.com/mit-han-lab)- StreamingLLM là một framework để triển khai LLM cho các ứng dụng streaming với đầu vào có độ dài vô hạn mà không phải đánh đổi hiệu quả hay hiệu năng
- Giải quyết vấn đề trong các cuộc hội thoại nhiều vòng cần tương tác dài, nơi việc cache trạng thái Key/Value của các token trước đó tiêu tốn nhiều bộ nhớ, và các LLM thông thường không thể tổng quát hóa tốt với văn bản dài hơn độ dài chuỗi đã dùng khi huấn luyện
- Window attention chỉ cache KV gần đây thường bị thất bại khi độ dài văn bản vượt quá kích thước cache, còn StreamingLLM khôi phục đáng kể hiệu năng của window attention bằng attention sink giữ lại KV của các token ban đầu
- Cho phép LLM được huấn luyện với attention window hữu hạn có thể tổng quát hóa sang độ dài chuỗi vô hạn mà không cần fine-tuning, đồng thời thực hiện language modeling ổn định và hiệu quả trên Llama-2, MPT, Falcon, Pythia với tối đa hơn 4 triệu token
- Trong thiết lập streaming, cho thấy tốc độ nhanh hơn tối đa 22.2 lần so với baseline sliding window recomputation
- Không mở rộng context window, chỉ giữ lại các token gần đây và attention sink, còn các token ở giữa sẽ bị loại bỏ
- Nếu Llama-2 được pretrain với context window 4096 token thì kích thước cache tối đa của Llama-2 trong StreamingLLM cũng là 4096
- Khi đưa vào một cuốn sách dài, mô hình chỉ nhận biết các token mới nhất nên chỉ có thể tóm tắt phần kết luận
- Trường hợp sử dụng phù hợp là các ứng dụng streaming cần hoạt động liên tục và phải tránh phụ thuộc vào dữ liệu quá khứ hoặc sử dụng nhiều bộ nhớ; ví dụ gồm hội thoại nhiều vòng và trợ lý hằng ngày dựa trên LLM
- Phương pháp này trực giao với các cách mở rộng ngữ cảnh gần đây và có thể tích hợp cùng nhau; trong ngữ cảnh của StreamingLLM, context extension nghĩa là có thể lưu nhiều token gần đây hơn nhờ kích thước cache lớn hơn
- Ví dụ chạy là
examples/run_streaming_llama.py --enable_streaming, và thiết lập môi trường sử dụng Python 3.8,torch,transformers==4.33.0,accelerate,datasets,evaluate,wandb,scikit-learn,scipy,sentencepiece - Mã nguồn cốt lõi cho Llama-2, MPT, Falcon, Pythia đã được công khai; mã đánh giá perplexity và demo Streaming Llama Chatbot cũng đã công khai, còn bộ dữ liệu StreamEval và mã đánh giá vẫn chưa được công bố
1 bình luận
Các ý kiến trên Hacker News
Có vẻ mọi người đang hiểu nhầm công trình này thành một thứ giống như attention dày đặc toàn phần
Điều được nói ở đây không phải là cho phép nhìn thấy nội dung ở rất xa, mà gần với một cải tiến hiệu quả hơn: tái sử dụng cache để duy trì perplexity, thay vì mỗi lần tính lại sliding window và chịu chi phí L² thêm T lần
Phép thử cũng được đo bằng cách nối kiểu
Q A Q A Q A Q A..., chứ không phải bắt mô hình tìm câu trả lời ở rất xa nhưQ Q Q Q A A A A...Đo perplexity có nghĩa là tạo ra “văn bản đọc được”, tức các câu có vẻ hợp lý về mặt cục bộ, chứ không phải bằng chứng rằng mô hình có thể “trích xuất” thứ gì đó từ một khoảng trống tam giác lớn mà attention không chạm tới
Nếu đưa cho nó một cuốn sách rồi yêu cầu viết từ đầu tiên của mỗi đoạn, hoặc tóm tắt mỗi chương thành một câu, tôi nghĩ nó sẽ thất bại
Tôi thử trực tiếp thì có vẻ nó không mở rộng độ dài ngữ cảnh, còn tốc độ chạy thì khá nhanh
Nó dùng khoảng 35GB bộ nhớ A100 và mức sử dụng cố định trong suốt quá trình chạy
Tôi lấy một cuốn sách từ Project Gutenberg, chia theo từng đoạn, đưa vào từng đoạn một và bảo nó trả lời “okay” cho mỗi đoạn, rồi hỏi ở cuối; nó đã hoàn toàn hallucinate câu trả lời
Nhân tiện, trong khoảng 10 phút nghịch thử, tôi cũng không làm được để mô hình mặc định
lmsys/vicuna-13b-v1.3trả lời bằng tiếng Anh một cách ổn địnhhttps://gist.github.com/bluecoconut/9cae9e91fe3b1616ed650a96...
Dù vậy đây vẫn là một công trình thú vị, và điểm cốt lõi có vẻ là phát hiện ở Figure 2
Hai layer đầu cho thấy một mẫu cục bộ, chú ý nhiều hơn đến các token gần đây, nhưng sau khi đi qua các layer bên dưới, mô hình lại chú ý rất mạnh đến các token ban đầu trên mọi layer và head
Các tác giả gọi chúng là “attention sinks”, và cho rằng dù chúng không quan trọng về mặt ngữ nghĩa, vì Softmax buộc tổng điểm attention phải bằng 1 nên phần attention dư phải đi về đâu đó
Trong mô hình ngôn ngữ tự hồi quy, các token ban đầu hầu như được mọi token phía sau nhìn thấy, nên chúng dễ được học để đóng vai trò sink như vậy
StreamingLLM gần giống một “hack” để hiệu chỉnh hành vi kỳ lạ này khi cắt cửa sổ attention của LLM để dùng, và vì đây là một trường hợp cho thấy vết nứt trong việc dùng Softmax, tôi nghĩ nếu muốn LLM linh hoạt với độ dài ngữ cảnh thì có lẽ một hàm khác sẽ phù hợp hơn
Lướt qua lần đầu thì nó tốt đến mức khiến người ta nghi ngờ có thật không, nhưng chất lượng công trình có vẻ ổn và kỹ thuật thì đơn giản đến đáng ngạc nhiên
Ý tưởng là ở mỗi layer, chỉ áp dụng attention lên token đầu tiên và cửa sổ ngữ cảnh trượt, còn các token ở giữa thì bỏ qua
Điều này có vẻ có nghĩa là mỗi layer đẩy dần thông tin liên quan về phía sau chuỗi, để cửa sổ attention trượt ở phần cuối của layer trên cùng có thể nhìn thấy nó
Tuy nhiên, nếu phạm vi mà tất cả các cửa sổ trượt bao phủ không đủ để nối toàn bộ chuỗi, thì có thể không truyền được mọi thông tin quan trọng về phía trước
Ví dụ, khi mọi cửa sổ có cùng độ dài, nếu
độ sâu mô hình × độ dài cửa sổ < độ dài chuỗithì sẽ xuất hiện giới hạnCó vẻ điều này khả thi nhờ quan sát rằng Softmax phải làm cho tổng bằng 1
Nhìn nhanh thì mô hình có xu hướng dùng token đầu tiên như một placeholder khi nó không cần chú ý đến token trước đó
Lần đầu tôi thấy vấn đề này là trong một bài HN của Evan Miller, nói rằng việc buộc các attention head phải phân bổ toàn bộ attention cho các token trước là sai, và nên cho phép “không chú ý” bằng cách cộng 1 vào mẫu số của Softmax
Việc tận dụng quan sát này mà không cần huấn luyện lại là tốt, và tôi cũng tò mò nếu làm theo đề xuất của Evan thì mô hình sẽ thay đổi ra sao
[2] https://news.ycombinator.com/item?id=36851494
Họ huấn luyện mô hình với một sink token chuyên dụng mà mọi giá trị đều bằng 0, nhưng có vẻ kết luận là ngay cả vậy các token ban đầu khác vẫn được dùng làm sink, nên có sink token chuyên dụng thì tốt hơn
Theo tôi biết thì nó chưa từng thật sự cải thiện hiệu năng
Khi thao tác cửa sổ attention sau huấn luyện, Softmax + 1 có thể phù hợp hơn, nhưng tôi không biết đã có ai thử ở quy mô lớn chưa
Thêm bộ nhớ cache attention là một cách giải rất thú vị cho vấn đề này
Vài ngày trước cũng có một paper đưa ra quan sát liên quan trong Vision Transformer
Các mô hình Transformer có vẻ như chọn ra các token để lưu thông tin toàn cục, và dường như cần một dạng “token để suy nghĩ”
Nếu cung cấp các token cụ thể cho mục đích này thì hiệu năng cải thiện đôi chút, và trực quan hóa để giải thích cũng khá thú vị
[0] https://arxiv.org/pdf/2309.16588.pdf
Khi tinh chỉnh, có thể giữ cố định tham số của mô hình gốc và chỉ điều chỉnh các tham số đi vào/đi ra từ các unit cache “tuning” mới
Như vậy có thể thay thế hoặc dùng cùng lúc nhiều bộ unit tuning khác nhau
Kiểu như trộn các siêu prompt: unit tránh từ tục + unit thuật ngữ cụ thể + unit viết ngắn gọn
Nếu số tham số mới đủ ít, dù tốn thêm bộ nhớ, việc tuning nhanh và hiệu quả bằng tối ưu hóa bậc cao cũng có thể khả thi
Cũng có thể nghĩ đến cách tăng đồng thời độ dài chuỗi và số unit trong quá trình huấn luyện
Với chuỗi ngắn thì chỉ dùng vài unit, rồi khi tăng độ dài chuỗi huấn luyện thì thêm unit và tiếp tục huấn luyện
Có lẽ cũng có thể điều khiển việc mở rộng cache bằng phân tích hiệu năng hoặc gradient thay vì một lịch trình tùy ý
Các tác giả đã đăng FAQ, có thể giúp làm rõ phần nào sự nhầm lẫn: https://github.com/mit-han-lab/streaming-llm/blob/main/READM...
Với câu hỏi “Có thể đưa một văn bản dài như sách vào StreamingLLM để tóm tắt không?”, họ trả lời rằng tuy có thể đưa văn bản dài vào, nhưng mô hình chỉ nhận biết các token mới nhất, nên nếu đưa cả cuốn sách vào thì nó chỉ có thể tóm tắt các đoạn cuối, và có thể không hữu ích lắm
Nói cách khác, đây không phải là mở rộng cửa sổ ngữ cảnh của LLM hay tăng cường trí nhớ dài hạn; điểm mạnh của StreamingLLM là sinh văn bản trôi chảy từ các token gần đây mà không cần làm mới cache
Có thể tôi sai, nhưng có vẻ đây không phải là kỹ thuật giúp LLM tham chiếu đến nội dung vượt quá độ dài đã được huấn luyện như mọi người nghĩ
Có lẽ nó gần với vấn đề duy trì hiệu năng của mô hình trên văn bản dài, chính xác hơn là hiệu năng đối với nội dung vẫn nằm trong cửa sổ ngữ cảnh
Lời giải thích là mô hình học cách đặt một dạng “hành lý” vào attention của các token đầu văn bản, và khi thứ đó biến mất khỏi cửa sổ thì mô hình bị hỏng; nhưng tôi không chắc vì sao lại vậy
Nếu không phải là phần nhập lệnh, tôi nghĩ phần văn bản ở giữa cũng tốt chẳng kém phần văn bản đầu
Tôi tò mò các kỹ thuật cửa sổ trượt như thế này xử lý ra sao khi một lệnh bất ngờ chỉ xuất hiện ở cuối
Ví dụ, nếu đưa một cuốn sách vào mô hình và câu cuối là “hãy trả về số lần ký tự m xuất hiện trong đầu vào trước đó”, con người sẽ thở dài rồi đọc lại để đếm, nhưng LLM không có khả năng quay lại và đọc lại đầu vào
Trong ví dụ này, ngay cả khi bỏ qua hạn chế cố hữu của LLM trong việc đếm ký tự, để giải quyết thật sự thì có lẽ LLM phải có khả năng tùy ý chạy vòng lặp và nhảy trong văn bản
Tất nhiên như vậy sẽ sinh ra những vấn đề hoàn toàn mới, và có thể cần hẳn một kiến trúc mới
Tức là chuyển câu hỏi thành các từ khóa liên quan, tra cứu, rồi đọc lướt lại nội dung để tìm thông tin liên quan
Nếu có bước tiền xử lý cần thiết, một LLM có thể “đi nghiên cứu đủ kỹ rồi mới trả lời” sẽ cực kỳ mạnh
Trong khoảng 10.000 năm qua, chúng ta đã cải thiện các kỹ thuật quản lý tri thức để vượt quá dung lượng và thời gian của bộ não cá nhân, nên các mô hình ngôn ngữ cũng nên tận dụng cách nghiên cứu thực sự và việc tiêu hóa trước, chứ không chỉ là một tìm kiếm Bing đơn giản
Trí nhớ ngắn hạn không cần nhớ từng đoạn mã đã làm gì; khi đọc chỉ cần gắn thẻ rồi dựa vào một chỉ mục thẻ dùng chung có thể mở rộng
Tuy nhiên, càng nghĩ thì điều này càng giống tiền huấn luyện LLM thông thường, và chỉ mục tri thức trông như một khối trọng số LLM khổng lồ
Nó giống một lớp đặt bên trên hơn là thay đổi hành vi của chính LLM
Có thể tôi hiểu nhầm, nhưng trong trường hợp này chẳng phải trạng thái ẩn giải quyết vấn đề truy xuất sao?
Dù sao trước khi trả lời nó cũng phải hấp thụ toàn bộ đầu vào, nên việc lệnh nằm ở đầu hay cuối dường như không ảnh hưởng nhiều ngoài attention
Claude 100K thì yêu cầu người dùng đặt lệnh ở cuối
Hoặc có thể dùng một mô hình nhanh để kiểm tra xem có lệnh ở cuối không rồi đưa nó lên phía trước
Tôi không chắc các mô hình hiện tại có làm được việc này ngay cả với đầu vào ngắn hay không
Nói hơi đùa một chút, LLM đang rất nỗ lực phát minh lại RNN, và nếu được cung cấp công cụ thì cuối cùng có lẽ sẽ làm được vậy
Nhìn theo cách khác, các mô hình Transformer đang cố dự đoán phần nào của mạng RNN “đáng được giữ lại” khi có ràng buộc tài nguyên
Transformer hiện nay dùng một heuristic đơn giản, và kết quả này làm heuristic đó tốt hơn
Như nhiều bài toán NP-đầy đủ, dù không khớp hoàn hảo vẫn có thể có các xấp xỉ hữu ích, và Transformer cho thấy điều đó cũng khả thi trong mạng nơ-ron
Nó đã từng đứng ở khoảng giữa trên các bảng xếp hạng mã nguồn mở một thời gian, nên đây là một hướng tiếp cận khá chính đáng, chỉ là không thịnh hành mà thôi
[1]: https://huggingface.co/blog/rwkv
Lợi thế chính của Transformer so với RNN là song song hóa quá trình huấn luyện
RNN gặp vấn đề tiêu biến gradient trong huấn luyện, và cũng khó nâng mức sử dụng tổng thể nên cần batch lớn, vì thế khá khó xử lý
Sự tồn tại của các mô hình như RWKV cho thấy có thể có một tương lai huấn luyện như Transformer và suy luận như RNN
Liên quan đến chủ đề này, Giáo sư Han của MIT đang tổ chức một khóa học TinyML công khai
https://news.ycombinator.com/item?id=37620507
https://efficientml.ai