MemAgent: Przełom w przetwarzaniu długich kontekstów przez LLM

Zdolność dużych modeli językowych (LLM) do przetwarzania i rozumienia obszernych dokumentów pozostaje jednym z kluczowych wyzwań w dziedzinie sztucznej inteligencji. Nawet zaawansowane techniki, takie jak ekstrapolacja długości kontekstu czy rzadkie mechanizmy uwagi, często prowadzą do spadku wydajności i generują wysokie koszty obliczeniowe. Prace naukowców z ByteDance Seed i Tsinghua University, którzy opracowali MemAgent, wskazują na potencjalne rozwiązanie tego problemu, oferując mechanizm o liniowej złożoności i minimalnej utracie precyzji.

Ograniczenia istniejących rozwiązań

Obecnie stosowane podejścia do modelowania długich kontekstów można podzielić na trzy główne kategorie. Pierwsza to metody ekstrapolacji długości, takie jak NTK, PI czy YaRN, które rozszerzają okno kontekstowe poprzez manipulację osadzeniami pozycyjnymi. Niestety, często wiążą się one ze spadkiem wydajności i problemami ze skalowaniem. Druga kategoria to rzadkie i liniowe mechanizmy uwagi, które redukują złożoność obliczeniową do O(n), ale zazwyczaj wymagają retrenowania modelu od podstaw i polegają na z góry określonych wzorcach lub regułach.

Trzecie podejście, kompresja kontekstu, wykorzystuje moduły pamięci do kondensacji długich danych wejściowych. Choć obiecujące, często zakłóca standardowe generowanie tekstu i ma trudności z ekstrapolacją na nowe, nieprzewidziane scenariusze.

Istniejące rozwiązania nie są w stanie zapewnić jednocześnie trzech kluczowych cech: wsparcia dla dowolnej długości danych wejściowych, stałej dokładności oraz efektywnej, liniowej złożoności obliczeniowej. To właśnie te braki stały się impulsem do powstania MemAgent.

MemAgent: ludzkie podejście do pamięci

MemAgent czerpie inspirację ze sposobu, w jaki ludzki umysł selekcjonuje kluczowe informacje, ignorując szum. System przetwarza dane wejściowe jako strumień „dowodów”. Na każdym etapie MemAgent odczytuje fragment dokumentu i swoją wewnętrzną pamięć, nadpisując ją zaktualizowanym, skompresowanym kontekstem. To podejście różni się od tradycyjnych, „ślepych” operacji na całym tekście.

Kluczowe innowacje MemAgent obejmują:

• Pamięć oparta na tokenach o stałej długości: Pozwala na kompresję kluczowych informacji przy jednoczesnym zachowaniu kompatybilności z modelem.
• Mechanizm nadpisywania segmentowego: Umożliwia przetwarzanie nieskończonych długości tekstu bez niekontrolowanego wzrostu pamięci.
• Liniowa złożoność: Koszt aktualizacji pamięci i dekodowania pozostaje stały dla każdego fragmentu danych.

Trening oparty na uczeniu wzmacniającym

MemAgent traktuje interakcję z każdym fragmentem dokumentu jako niezależny dialog. Model jest trenowany za pomocą metody GRPO (Group Relative Policy Optimization) w ramach potoku uczenia wzmacniającego wielu konwersacji, zwanego DAPO. Takie podejście umożliwia aktualizację pamięci sterowaną przez nagrody, co jest ewolucją w stosunku do tradycyjnych, mniej „inteligentnych” metod.

Istotnymi elementami tego systemu są weryfikator oparty na regułach, który oblicza nagrody wynikowe poprzez porównanie odpowiedzi modelu z wieloma referencyjnymi odpowiedziami, oraz sygnał uczenia wzmacniającego na poziomie tokenów, stosowany równomiernie w konwersacjach. To wszystko zachęca do kompresji pamięci, skupionej na informacjach istotnych dla odpowiedzi, jednocześnie odrzucając zbędne dane.

Ocena wydajności i studium przypadku

W testach MemAgent, trenowany z 8K okna kontekstowego, wykazał zdolność do ekstrapolacji aż do 3,5 miliona tokenów. Model utrzymywał ponad 95% dokładności w benchmarkach RULER (od 8K do 512K tokenów) i konsekwentnie przewyższał baseliny oparte na długich kontekstach i destylacji.

W studium przypadku dotyczącym wieloetapowej odpowiedzi na pytania (Multi-Hop QA), MemAgent poprawnie przetworzył zapytanie „Reżyser romantycznej komedii ‘Big Stone Gap’ ma siedzibę w jakim mieście w Nowym Jorku?” Model progresywnie śledził istotne treści. Początkowo rozpoznawał nieistotne fragmenty, ale zachowywał informacje o lokalizacji. Następnie utrzymywał pamięć pomimo kolejnych nieistotnych części, by w końcu precyzyjnie zaktualizować swoją pamięć po napotkaniu biografii Adriany Trigiani. Końcowa odpowiedź była poprawna: Greenwich Village, Nowy Jork.

Fundamenty teoretyczne i złożoność

MemAgent reformuluje model autoregresywny, używając zmiennych pamięci latentnej, co umożliwia koszt obliczeniowy O(N) i czytelną dla człowieka pamięć pośrednią — w przeciwieństwie do kompresji cech opartej na uwadze. Uczenie wzmacniające jest tu kluczowe, ponieważ aktualizacje pamięci są dyskretne i nie mogą być nauczane poprzez samo wsteczne propagowanie błędu.

W rezultacie MemAgent oferuje skalowalne i efektywne rozwiązanie tzw. trylematu długiego kontekstu: nieograniczonej długości danych wejściowych, niemal bezstratnej dokładności i liniowej złożoności. Mechanizm nadpisywania pamięci oparty na uczeniu wzmacniającym pozwala modelom LLM na czytanie, abstrakcję i generowanie tekstu na podstawie danych wejściowych liczących miliony tokenów, bez konieczności modyfikacji architektury.