🌟 NVIDIA переизобретает память: LLM, которые доучиваются прямо во время инференса

Контекстные окна растут, но тут два стула: либо классическое внимание, которое питается памятью и компьютит как не в себя, либо RNN-подобные Mamba, DeltaNet, которые работают быстро, но в длинном контексте начинают плыть и терять детали.

NVIDIA предлагает решение, которое пытается усидеть на обоих стульях сразу - Test-Time Training with End-to-End formulation (TTT-E2E):

Обычно веса модели заморожены после тренировки. Когда вы скармливаете ей данные, она просто держит её в KV-кэше. В TTT все по-другому: контекст — это и есть обучающий датасет. Пока модель читает ваш промпт (контекст), она обновляет свои веса (если точнее - делает градиентный спуск прямо на лету), тем самым, инфа из контекста впекается в саму модель. Это позволяет сжать гигантские объемы в фиксированный размер состояния, не раздувая KV-кэш до небес.

🟡В результате - красота и волшебство:

🟢Латентность инференса становится константной. Неважно, 100 токенов в контексте или миллион — время генерации следующего токена одинаковое.

🟢На контексте 128k токенов — ускорение в 2.7x по сравнению с Attention (на H100). На 2M токенов — ускорение в 35 раз.

🟢В отличие от Mamba и других RNN, качество не проседает на длинных дистанциях. TTT держит планку на уровне полного внимания.

🟡Разумеется, есть куча пунктов со звездочкой

🟠Трейн - сложный. Чтобы модель могла так лихо учиться на лету, её нужно спечиальным образом претрейнить. Этот процесс сейчас в 3.4x медленнее, чем обычное обучение.

🟠Метод требует вычисления градиентов от градиентов во время обучения. FlashAttention из коробки сейчас это не поддерживает, нужны кастомные ядра или костыли.

🟠Cам процесс поедания контекста при инференсе требует вычислений во время префилл-фазы.

По итогу, NVIDIA сравнивает RAG с блокнотом, а свой TTT — с реальным обновлением нейронных связей мозга. Если есть желание покопаться в методике и проникнуться идеей - код и пейпер в открытом доступе.


🟡Статья
🟡Arxiv
🖥GitHub


@ai_machinelearning_big_data

#AI #ML #LLM #TTTE2E #NVIDIA