🐹 ARC-Encoder от Kyutai: plug-and-play компрессия для LLM

На прошлой неделе вышла крутая статья от Kyutai Labs (да, я им симпатизирую): 🔗ARC-Encoder learning compressed text representations for large language models

Почему-то осталась относительно незамеченной, хотя на мой взгляд работа хорошо сделана и достойна deep-dive чтения. Как уже видно из названия, эта статья даже не про TTS & STT, здесь авторы придумали универсальный энкодер для сжатия длинных контекстов в LLM'ках в х4-8 раз практически без потерь качества ответов на различных downstream задачах (QA, text summarization, перевод), включая ICL

ARC-Encoder сжимает контекст в короткую последовательность эмбеддингов, которые напрямую подставляются в LLM декодер (после embedding layer), при этому саму LLM не до-обучают❕. И в дополнение, ARC-Encoder можно легко адаптировать к нескольким разным декодерам через небольшие MLP модули

🍅 Немного про context compression

LLM-пайплайны (RAG, multi-document QA) нуждаются в длинных контекстах. Но pre-fill на длинных промптах стоит дорого и быстро упирается в лимиты контекстного окна. Можно просто дропнуть наименее важные токены, или кратко суммаризировать промпт. Это hard compression. Альтернативный подход, soft compression, подразумевает извлечение смысловых эмбеддингов из пассажей контекста (инструкции промпта или документы). Часто такие методы также подразумевают fine-tune LLM декодера, ведь ему нужно еще научиться использовать эти сжатые репрезентации. ARC-Encoder же тем хорош, что никак не трогает decoder LLM, таким образом модель никак не будет менять своего поведения на других задачах 💲🤩


☕️ Архитектура

⭐️В недавней работе (тут) показали, что LLM декодер можно успешно использовать для компресси текста. Поэтому авторы взяли за основу LLaMA‑3.2‑3B, убрав output head и cаusal mask (чтобы энкодер видел весь контекст сразу, а не только влево)

⭐️Поверх последнего self-attention блока применяют pooling (average) с фактором сжатия от 4 до 8, но только для queries, в то время keys & values остаются полноразмерными. Авторы пишут, что пробовали пулить на разных слоях, но эффективнее всего на последнем, что логично, репрезентации на последнем слое семантически богаче

⭐️Далее MLP слой (без активации🌸) отображает encoder outputs в размерность латентного пространства LLM. Получившиеся эмбеддинги просто подставляются в латентное пространство LLM декодера (на вход transformer layers)

Как и что тренили?

1️⃣ Pre-training на Common Crawl на двух задачах (см. картинку): реконструкция последовательности (а) и продолжение последовательности (б), что оказалось критично, так как это по сути та же задача, что и любой downstream task

2️⃣ Fine-tuning делали на миксе синтетических и реальных QA, summarization и др. датасетах. При этом сжатые эмбеддинги перемежаются с обычными токенами, что сохраняет few-shot способности модели

Понравился оценочный сетап. Модели‑декодеры: Llama‑3.1‑8B и Mistral‑7B (base). Бейзлайны авторы ре-имплементили сами на базе тех же декодеров, данных и задачах. Дополнительно оценили open book (весь промпт вместе с документами для QA у модели в контексте) и closed book сетапы. Тренировочные сеты бенчмарков предусмотрительно не включали в тренировки.

🟣На всех бенчах context compression ARC держится близко к open‑book на Llama‑3.1‑8B и обгоняет closed‑book и большинство бейзлайнов; на Mistral‑7B отставание от open‑book небольшое, что ожидаемо (архитектурно декодер Mistral дальше от Llama-3.2 энкодера)

🔴Также круто: удалось увеличить эффективнй контекст Llama2-chat с 4 до 32k обогнав по некоторым метрикам Llama2-32k Instruct (специально до-трененную на длинных последовательностях)

🥹 А что же дальше?

На мой взгляд вся идея хорошо ложится на речевые и аудио модели, в которых проблема длинных контекстов стоит еще острее. Как думаете, является ли эта работа своего рода proof of concept и увидим ли мы эту идею в следующем Moshi? Также интересно, как далеко можно зайти с multi-decoder подходом: один универсальный энкодер для сжатия + мини адаптеры под каждую LLM