2025-11-16T06:22:12.451775

To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models

Malach, Saremi, Williamson et al.
State Space Models (SSMs) have become the leading alternative to Transformers for sequence modeling. Their primary advantage is efficiency in long-context and long-form generation, enabled by fixed-size memory and linear scaling of computational complexity. We begin this work by showing a simple theoretical result stating that SSMs cannot accurately solve any ``truly long-form'' generation problem (in a sense we formally define), undermining their main competitive advantage. However, we show that this limitation can be mitigated by allowing SSMs interactive access to external tools. In fact, we show that given the right choice of tool access and problem-dependent training data, SSMs can learn to solve any tractable problem and generalize to arbitrary problem length/complexity (i.e., achieve length generalization). Following our theoretical finding, we demonstrate that tool-augmented SSMs achieve remarkable length generalization on a variety of arithmetic, reasoning, and coding tasks. These findings highlight SSMs as a potential efficient alternative to Transformers in interactive tool-based and agentic settings.
academic

К бесконечности и дальше: использование инструментов разблокирует обобщение длины в моделях пространства состояний

Основная информация

  • ID статьи: 2510.14826
  • Название: To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models
  • Авторы: Eran Malach, Omid Saremi, Sinead Williamson, Arwen Bradley, Aryo Lotfi, Emmanuel Abbe, Josh Susskind, Etai Littwin
  • Учреждение: Apple
  • Классификация: cs.LG
  • Дата публикации: 17 октября 2025 г.
  • Ссылка на статью: https://arxiv.org/abs/2510.14826

Аннотация

Модели пространства состояний (SSMs) стали основной альтернативой Трансформерам в моделировании последовательностей, основное преимущество которых заключается в эффективности обработки длинного контекста и генерации длинных последовательностей благодаря фиксированному размеру памяти и линейной вычислительной сложности. В данной работе сначала представлен простой теоретический результат, доказывающий, что SSMs не могут точно решить никакую "истинно длинную" задачу генерации последовательностей (в формально определённом смысле), что ослабляет их основное конкурентное преимущество. Однако исследование показывает, что это ограничение может быть преодолено путём предоставления SSMs интерактивного доступа к внешним инструментам. На практике, при правильном выборе доступа к инструментам и соответствующих данных обучения, SSMs могут научиться решать любую разрешимую задачу и обобщаться на произвольные длины/сложность задач. На основе теоретических результатов авторы демонстрируют, что дополненные инструментами SSMs достигают значительного обобщения длины на различных арифметических, логических и программистских задачах.

Исследовательский контекст и мотивация

Постановка проблемы

  1. Вычислительные узкие места Трансформеров: Трансформеры из-за механизма внимания имеют квадратичный рост вычислительной сложности с длиной последовательности и линейный рост памяти, что становится основным ограничением при обработке длинного контекста и генерации длинных последовательностей.
  2. Возникновение SSMs: Для решения этой проблемы исследователи предложили различные альтернативные архитектуры, такие как линейные Трансформеры и модели пространства состояний (SSMs), включая Mamba, DeltaNet и другие, которые достигают фиксированной памяти и линейной вычислительной сложности.
  3. Ограничения SSMs: Несмотря на преимущества в эффективности, некоторые исследования указывают на значительные ограничения SSMs в задачах, требующих долгосрочной памяти последовательности и контекстного обучения.

Исследовательская мотивация

Авторы стремятся понять возможности и ограничения SSMs в задачах генерации длинных последовательностей, особенно в тех, где выходная длина растёт с увеличением сложности задачи. Это именно те типы задач, где SSMs демонстрируют явные преимущества в эффективности рассуждений по сравнению с Трансформерами.

Основные вклады

  1. Теоретический отрицательный результат: Доказано, что SSMs не могут точно решить "истинно длинные" задачи генерации последовательностей, даже при разрешении произвольно длинной генерации цепочки мыслей (CoT).
  2. Теоретическая база для использования инструментов: Введена новая теоретическая база для изучения агентов ReAct, доказывающая, что интерактивное использование инструментов может значительно расширить возможности SSMs.
  3. Теорема о достаточности обобщения длины: Доказано, что SSMs, оснащённые надлежащим доступом к инструментам и специфическими данными обучения, могут достичь обобщения длины на любой разрешимой задаче генерации длинных последовательностей.
  4. Экспериментальная верификация: Продемонстрирована превосходная способность обобщения длины дополненных инструментами SSMs на арифметических, логических и программистских задачах.

Подробное описание методологии

Определение задач

Формальное определение задачи генерации длинных последовательностей:

  • Пусть Σ — словарь, X₁,X₂,... и Y₁,Y₂,... — последовательности входных и выходных пространств соответственно
  • D₁,D₂,... — последовательность распределений, где Dₙ — распределение на Xₙ
  • f: Σ* → Σ* — истинная функция, удовлетворяющая f(Xₙ) ⊆ Yₙ

Определение 2.2: Пара (f, {Dₙ}) называется задачей генерации длинных последовательностей с покрытием α тогда и только тогда, когда suppₐ(f(Dₙ)) монотонно возрастает с n и limₙ→∞ suppₐ(f(Dₙ)) = ∞.

Обобщённая модель пространства состояний (GSSM)

Определение: GSSM определяется следующими компонентами:

  • Пространство состояний S (конечное множество)
  • Начальное состояние s₀ ∈ S
  • Правило обновления u: S × Σ → S
  • Правило вывода r: S → Δ(Σ)

Параметры использования инструментов:

  1. Только CoT: Разрешены только токены размышления и вывода
  2. Однораундовое использование инструментов: Разрешён единственный вызов инструмента
  3. Интерактивное использование инструментов: Разрешены произвольные вызовы инструментов и свободное чередование

Основные теоретические результаты

Теорема 2.1 (отрицательный результат): Для любой задачи генерации длинных последовательностей f с покрытием α существует сложность задачи n₀, такая что для всех n ≥ n₀ любой GSSM h, использующий только CoT или однораундовое использование инструментов, имеет частоту ошибок: errₙ(h) ≥ 1-α.

Теорема 2.2 (положительный результат): Существуют оракул инструмента памяти O и простой алгоритм обучения GSSM A, такие что для любой вычислимой задачи генерации длинных последовательностей f существует последовательность распределений обучения {Pₙ}, при которых A достигает обобщения длины в интерактивной установке.

Технические инновации

  1. Проектирование инструмента памяти: Предоставление доступа к внешней памяти с операциями чтения/записи через указатели, позволяющие моделировать операции машины Тьюринга.
  2. Интерактивная парадигма обучения: Путём конструирования данных обучения, содержащих траектории использования инструментов, SSMs учатся использовать внешнюю память для преодоления ограничений внутренней памяти.
  3. Генерация алгоритмических траекторий: Разработка синтетических траекторий использования инструментов для различных задач (сложение, умножение, логические рассуждения и т.д.), точно моделирующих требуемые алгоритмы.

Экспериментальная установка

Наборы данных

  1. Арифметические задачи: Многозначное сложение и умножение, обучение до 5-10 цифр, тестирование до 1000 цифр
  2. Ханойские башни: Обучение до 8 дисков, тестирование до 12 дисков
  3. Логические рассуждения на графах: Обучение до 10 узлов, тестирование до 1000 узлов
  4. Исправление кода: Обучение на кодовых базах до 16 функций, тестирование на больших масштабах

Конфигурация моделей

  • SSMs: Mamba-130M/1.4B, LSTM, GRU
  • Трансформеры: Pythia-160M/1.4B, Mistral (скользящее окно внимания)
  • Все модели сопоставимого размера (~130M параметров)

Типы инструментов

  1. Память с указателями: Поддержка инициализации, перемещения, операций чтения
  2. Инструменты поиска: Поддержка поиска шаблонов в контексте
  3. Команды Bash: Операции с файлами для задач исправления кода

Результаты экспериментов

Основные результаты

Производительность на арифметических задачах:

  • Mamba может идеально выполнять сложение 1000-значных чисел после обучения на 5-значных (100% точность)
  • Задачи умножения: обучение 10×1 цифр → тестирование 1000×1 цифр (100% точность)
  • Модели Трансформеров практически не обобщаются за пределы длины обучения

Производительность на задачах рассуждений:

  • Логические рассуждения на графах: обучение 10 узлов → тестирование 1000 узлов (98% точность)
  • Ханойские башни: обучение 8 дисков → тестирование 12 дисков (49% точность при экспоненциальном росте выходной длины)

Исправление кода:

  • При интерактивном обучении агентов Mamba сохраняет лучшую производительность на больших кодовых базах
  • Трансформеры показывают лучшую производительность на малых кодовых базах, но не обобщаются на большие масштабы

Абляционные исследования

Ключевые находки:

  1. Удаление CoT или использования инструментов приводит к почти полной потере способности обобщения длины
  2. Однораундовое использование инструментов имеет ограниченный эффект, интерактивное использование критично
  3. Смешанное обучение на задачах может улучшить обобщение при ограниченном бюджете

Экспериментальные находки

  1. Преимущества архитектуры: SSMs/RNNs значительно превосходят Трансформеры в установке с дополненными инструментами
  2. Важность интерактивности: Интерактивное использование инструментов является ключом к достижению обобщения длины
  3. Качество данных обучения: Тщательно сконструированные алгоритмические траектории критичны для успеха
  4. Масштабируемость: Метод масштабируется на различные алгоритмические задачи

Связанные работы

Основные направления исследований

  1. Цепочка мыслей и черновики: CoT значительно улучшает способность рассуждений LLMs, теоретически улучшая выразительность и обучаемость
  2. Нейронные машины Тьюринга: Ранние попытки моделирования машин Тьюринга нейронными сетями, но не получившие широкого распространения
  3. Обобщение длины: Обширные работы по изучению обобщения длины Трансформеров с предложением различных методов улучшения

Вклад данной работы

  • Первое систематическое исследование теоретических ограничений обобщения длины SSMs
  • Предложение использования инструментов как эффективного способа преодоления ограничений
  • Анализ производительности архитектур в контексте агентских систем, а не независимых моделей

Заключение и обсуждение

Основные выводы

  1. SSMs при независимом использовании имеют фундаментальные ограничения обобщения длины
  2. Интерактивное использование инструментов может полностью преодолеть эти ограничения
  3. В агентских установках SSMs могут превосходить Трансформеры

Ограничения

  1. Алгоритмы обучения в теоретическом анализе относительно просты (сопоставление строк)
  2. Обобщение ограничено для задач с экспоненциальной выходной длиной, таких как Ханойские башни
  3. Требуется тщательное проектирование траекторий обучения
  4. Ограниченное обобщение в задачах исправления кода

Будущие направления

  1. Разработка большего количества агентов на основе SSM с использованием инструментов
  2. Исследование теоретических гарантий для более естественных алгоритмов обучения (например, градиентный спуск)
  3. Расширение на более сложные задачи рассуждений и агентские системы
  4. Исследование потенциала гибридных архитектур

Глубокая оценка

Преимущества

  1. Теоретическая строгость: Предоставлены строгие математические доказательства ограничений SSMs
  2. Практическая ценность: Продемонстрирована реальная эффективность использования инструментов
  3. Полнота экспериментов: Охватывают различные типы задач и архитектуры моделей
  4. Глубокие инсайты: Раскрывают, что производительность архитектур в системах может отличаться от независимого использования

Недостатки

  1. Разрыв между теорией и практикой: Простые алгоритмы обучения в теоретическом анализе отличаются от фактического обучения нейронных сетей
  2. Ограничения задач: Основное внимание на алгоритмических задачах, применимость к открытой генерации неясна
  3. Инженерная сложность: Требуется проектирование специфических инструментов и траекторий обучения для каждой задачи
  4. Проблемы масштабируемости: Производительность на более сложных реальных задачах ещё требует проверки

Влияние

  1. Теоретический вклад: Предоставляет новую перспективу для понимания фундаментальных различий в возможностях различных архитектур
  2. Практическое руководство: Обеспечивает теоретическую поддержку применения SSM в агентских системах
  3. Направление исследований: Может стимулировать дальнейшие исследования дополненных инструментами языковых моделей

Применимые сценарии

  1. Алгоритмическое выполнение: Задачи, требующие точного выполнения известных алгоритмов
  2. Обработка длинных последовательностей: Сценарии с ограниченными вычислительными ресурсами, но требующие обработки длинных последовательностей
  3. Агентские системы: Приложения интеллектуальных агентов, требующие взаимодействия с внешними инструментами
  4. Образовательные приложения: Системы обучения, демонстрирующие процесс выполнения алгоритмов

Библиография

Данная работа ссылается на важные работы в этой области, включая:

  • Оригинальную статью о Трансформерах (Vaswani et al., 2017)
  • Архитектуры SSM, такие как Mamba (Gu & Dao, 2023)
  • Исследования цепочки мыслей (Wei et al., 2022)
  • Фреймворк ReAct (Yao et al., 2023)
  • Работы по обобщению длины (Zhou et al., 2024 и др.)

Резюме: Это высококачественная статья, сочетающая теорию и эксперименты, предоставляющая важные инсайты для понимания границ возможностей SSMs и ценности использования инструментов. Хотя масштабируемость в практических приложениях ещё требует проверки, её теоретические вклады и экспериментальные результаты имеют важное значение для развития этой области.