Google представи Titans и MIRAS – следващото поколение AI за свръхдълги контексти

Google представи нова архитектура за работа с дълги последователности Titans и теоретичната основа MIRAS, която съчетава скоростта на рекурентните мрежи с точността на трансформърите. Компанията заяви, че този подход помага на моделите да запомнят по-добре важните детайли в много дълги контексти – от големи документи до геномни данни.

Класическите трансформъри революционизират механизма за внимание: моделът може да „погледне назад“ и да подчертае значими части от текста. Но това си има цена.Изчислителната сложност на вниманието нараства квадратично с дължината на последователността, така че мащабирането на трансформъри за контексти с милиони лексеми става твърде ресурсоемко.

Изследователската общност вече се опита да заобиколи това ограничение чрез по-бързи линейни архитектури – например ефективните RNN и моделите от клас State Space като Mamba-2. Те компресират предишния контекст във вектор с фиксиран размер, така че изчислителните разходи нарастват линейно. Но именно поради сериозното ограничение на размера на паметта тези модели не се справят добре с наистина богатите и дълги последователности, при които е важно да не се губят фините връзки между фактите.

В два нови документа Google предлага съвкупност от Titans и MIRAS. Titans е конкретна архитектура, а MIRAS е обща теоретична рамка, която описва подобни системи като вид асоциативна памет. Заедно те развиват идеята за „инференциална памет“: даден модел може да се доучава непосредствено по време на изпълнението му, не само чрез извличане на информация от параметрите, но и чрез актуализиране на дългосрочната памет при наличието на данни.

Основната идея зад Titans е, че моделът има няколко нива на памет, подобно на краткосрочната и дългосрочната памет при хората. Краткосрочната част все още се осигурява от механизма за внимание, който работи добре върху локализирани участъци от текста. Но дългосрочната памет е реализирана не като обикновен вектор или матрица с фиксиран размер, а като пълноценна дълбока невронна мрежа – многослоен перцептрон. Благодарение на това моделът може да кодира миналия контекст много по-богато и да „разбира“ историята, а не просто да съхранява набор от бележки.

Друга ключова идея е механизмът „изненада“. В живота лесно забравяме рутинните действия, но отлично запомняме събитията, които излизат от рамките на даден модел. Titans използва подобен принцип: моделът сравнява текущото състояние на паметта с новия входящ сигнал и изчислява колко „неочакван“ е той за нея. Ако несъответствието е малко, информацията може да не бъде записана в дългосрочната памет. Ако входният сигнал рязко се различава от очакванията, това е сигурен знак, че разполагаме с важен или аномален факт, който си струва да бъде запазен.

За да се гарантира, че системата не реагира само на еднократни скокове, в Titans са добавени два елемента. На първо място, „инерцията“ – моделът взема предвид не само настоящата изненада, но и скорошната история, за да улови вериги от свързани събития. Второ, механизъм за забравяне чрез адаптивно намаляване на теглото. Това е своеобразен „клапан“, който предпазва паметта от препълване при работа с изключително дълги последователности и постепенно изхвърля остарялата информация.

MIRAS описва целия този клас модели на по-абстрактно ниво. Тази теория разглежда всяка архитектура на последователности като асоциативна памет, която се научава да съпоставя ключове и стойности, като балансира двете задачи – да научава новото и да не унищожава старото. MIRAS подчертава четири ключови компонента: структурата на самата памет, как моделът избира на какво да обърне внимание, механизма за запазване и забравяне и оптимизационния алгоритъм, чрез който паметта се обновява.

Google поставя отделен акцент върху факта, че MIRAS излиза извън рамките на обичайните евклидови метрики като средна квадратична грешка и скаларно произведение. В повечето съвременни модели именно MSE и точковото произведение са в основата както на „пристрастията“ на модела, така и на механизмите за регуларизация на паметта. Това е удобно, но прави системата чувствителна към външни отклонения и ограничава класа на приемливите решения. MIRAS предлага по-богато пространство от цели и регулатори, като се опира на прозрения от оптимизацията и статистиката, и по този начин отваря пътя към архитектури с нестандартни, неевклидови функции на загубите.

В тази рамка изследователите на Google са разработили три специфични модела без явно внимание – YAAD, MONETA и MEMORA. YAAD използва по-меко наказание за грешките чрез загубите на Хубер и по-добре понася зашумените данни, при които са възможни печатни грешки или единични артефакти. MONETA експериментира с по-сложни норми и проверява дали строгите математически „правила“ за запомняне и забравяне могат да подобрят стабилността на дългосрочната памет. MEMORA се опитва да стабилизира максимално паметта, като я кара да се държи като вероятностно разпределение, така че всяко обновяване на състоянието да е контролирано и точно.

В експериментите моделите Titans и MIRAS бяха сравнени с най-съвременните архитектури като Transformer++, Mamba-2 и Gated DeltaNet. Тестването е извършено върху стандартни набори от данни за езиково моделиране, като C4 и WikiText, както и върху задачи с нулева стъпка на здравия разум, включително HellaSwag и PIQA. Google твърди, че новите модели показват по-ниска перплексия и по-добра точност от базовите архитектури със сравним размер, като същевременно запазват линейната мащабируемост и възможностите за паралелно обучение.

Предимството на Titans е особено забележимо при задачите с изключително дълги контексти. При бенчмарка BABILong, където е необходимо да се направят заключения въз основа на факти, разпръснати в много дълги документи, архитектурата Titans превъзхожда всички базови модели, включително големи системи като GPT-4, въпреки че съдържа много по-малко параметри. Авторите също така отбелязват, че Titans успява да мащабира контекстния прозорец до стойности над 2 милиона токена и все пак да запази качеството.

Отделни изследвания показват, че дълбочината на модула на паметта играе решаваща роля. За един и същ „размер“ на паметта, но с различна дълбочина, по-дълбоките варианти дават последователно по-ниска перплексия и се мащабират по-добре с увеличаване на дължината на последователността. Това е допълнителен аргумент, че богатата, дълбока дългосрочна памет наистина помага на моделите да се ориентират в много дългите последователности.

В резултат на това Google позиционира Titans и MIRAS като стъпка към ново поколение модели, които могат да се учат в движение, да съхраняват важни детайли върху огромни пространства от контекст и все пак да бъдат достатъчно ефективни за практическа употреба. Теоретичната рамка на MIRAS свързва онлайн оптимизацията, асоциативната памет и архитектурния дизайн, докато самата архитектура Titans демонстрира как скоростта на RNN може да се съчетае с изразителната сила на Transformers в ерата на изкуствения интелект с дълъг контекст.