Модели пространства состояний для табличных сетей Prior-Data Fitted Networks

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

• ID статьи: 2510.14573
• Название: State-Space Models for Tabular Prior-Data Fitted Networks
• Авторы: Felix Koch, Marcel Wever, Fabian Raisch, Benjamin Tischler
• Классификация: cs.LG
• Дата публикации/конференция: Proceedings of the 1st ICML Workshop on Foundation Models for Structured Data, Vancouver, Canada. 2025
• Ссылка на статью: https://arxiv.org/abs/2510.14573

Аннотация

Недавние достижения в области фундаментальных моделей для табличных данных, таких как TabPFN, продемонстрировали, что предварительно обученные архитектуры Transformer могут аппроксимировать байесовский вывод с высокой предсказательной производительностью. Однако Transformers страдают от квадратичной сложности относительно длины последовательности, что мотивирует исследование более эффективных моделей последовательностей. В данной работе мы исследуем потенциал использования Hydra, двунаправленной линейной модели структурированного пространства состояний (SSM), в качестве альтернативы Transformers в TabPFN. Ключевая проблема заключается в врождённой чувствительности SSM к порядку входных токенов — нежелательное свойство для табличных наборов данных, где порядок строк семантически бессмыслен. Мы исследуем, в какой степени двунаправленный подход может сохранить эффективность и обеспечить симметричную агрегацию контекста. Наши эксперименты показывают, что этот подход снижает зависимость от порядка, достигая предсказательной производительности, конкурентоспособной с исходной моделью TabPFN.

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

1. Решаемая проблема: Данное исследование направлено на решение проблемы вычислительной эффективности архитектуры Transformer в фундаментальных моделях для табличных данных, в частности её O(n²) сложность ограничивает масштабируемость на больших наборах данных.
2. Важность проблемы: TabPFN как фундаментальная модель для табличных данных демонстрирует отличную производительность, способна аппроксимировать байесовский вывод на миллисекундном уровне, однако её архитектура на основе Transformer сталкивается с узкими местами памяти и вычислений при обработке крупномасштабных данных.
3. Ограничения существующих методов:
   • Механизм самовнимания Transformer имеет квадратичную сложность
   • Прямая замена Transformer на Mamba вводит чувствительность к порядку входной последовательности
   • Порядок строк в табличных данных семантически бессмыслен, что противоречит причинному дизайну SSM
4. Исследовательская мотивация: Исследование структурированных моделей пространства состояний (SSM) как альтернативы Transformer, сохраняя преимущества линейной сложности, одновременно снижая зависимость от порядка входных данных через двунаправленный механизм обработки.

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

1. Предложена архитектура TabPFN на основе Hydra: Интеграция двунаправленной структурированной модели пространства состояний Hydra в TabPFN, обеспечивающая линейную временную сложность обработки табличных данных.
2. Введена техника повторяющейся перестановки контекста (RCP): Дальнейшее снижение чувствительности SSM к порядку последовательности путём многократной случайной перестановки входных данных и усреднения результатов предсказания.
3. Достигнут значительный прирост масштабируемости: По сравнению с исходной TabPFN новый метод может обрабатывать наборы данных на два порядка больше (расширение с 2¹⁵ до 2¹⁷ строк).
4. Сохранена конкурентоспособная предсказательная производительность: На эталонном тесте OpenML CC-18 точность Hydra-based TabPFN всего на 1,1% ниже исходной модели.

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

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

Данная работа исследует задачу табличной классификации, где:

• Входные данные: Полный набор табличных данных, содержащий обучающие и тестовые образцы
• Выходные данные: Предсказание вероятностей класса для тестовых образцов
• Ограничения: Требуется завершить вывод в одном прямом проходе без обновления градиентов или дополнительной настройки

Архитектура модели

1. Замена архитектуры Hydra

• Основной дизайн: Замена стека слоёв кодировщика Transformer на слои Hydra
• Двунаправленная обработка: Использование квазиразделимых матричных миксеров для двунаправленного моделирования пространства состояний
• Структура слоя: Каждый слой Hydra содержит двунаправленное смешивание пространства состояний, за которым следует преобразование прямого распространения

2. Сохранение стратегии встраивания

• Сохранение исходного метода встраивания данных TabPFN
• Каждый вход представляется как конкатенация значений признаков и меток класса
• При выводе немеченые данные обрабатываются путём маргинализации всех возможных назначений меток

3. Повторяющаяся перестановка контекста (RCP)

Алгоритм работает следующим образом:

Вход: количество перестановок r, контекст D, тестовый образец xtest
Выход: предсказанное значение класса
Инициализация пустого списка: outputs ← []
для i = 1 до r выполнить
    Перемешать строки D: Dp ← shuffle(D)
    Конкатенировать xtest с Dp: Din ← Dp ∪ xtest
    Предсказание: outputs[i] ← PFN.predict(Din)
конец цикла
Вернуть среднее значение outputs

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

1. Двунаправленность решает чувствительность к порядку: По сравнению с однонаправленной Mamba, двунаправленная обработка Hydra может симметрично агрегировать информацию контекста, снижая зависимость от порядка входных данных.
2. Линейная сложность: Реализация O(n) сложности через квазиразделимые матричные умножители, что даёт значительное преимущество перед O(n²) Transformer.
3. Стратегия RCP: Инновационное снижение чувствительности к порядку путём многократной случайной перестановки и усреднения результатов, что является специализированным дизайном для характеристик табличных данных.

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

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

• Основной набор данных: Эталонный набор тестов OpenML CC-18
• Критерии фильтрации: ≤2000 строк, ≤100 признаков, ≤10 классов
• Итоговые наборы данных: 30 наборов данных многоклассовой классификации
• Разделение данных: Каждый набор данных случайно разделён на обучающий/тестовый наборы 16 раз

Метрики оценки

1. Точность (Accuracy): Доля правильно классифицированных образцов
2. AUC OvO: Многоклассовая AUC "один против одного"
3. Дивергенция KL: Измерение различий в распределении предсказаний при различных порядках входных данных, оценка чувствительности к порядку
4. Время вывода: Время вычисления при различных размерах входных данных
5. Использование памяти: Максимальный размер набора данных, который может быть обработан

Методы сравнения

• TabPFN на основе Transformer: Исходная базовая модель
• TabPFN на основе Mamba: Однонаправленная замена SSM
• TabPFN на основе Hydra: Предложенный двунаправленный вариант SSM

Детали реализации

• Оборудование для обучения: Nvidia A40 GPU (48GB)
• Оборудование для тестирования: NVIDIA H100 80GB
• Время обучения: Transformer 48 часов, Mamba 52 часа, Hydra 134 часа
• Ключевые гиперпараметры:
  • Скорость обучения: 0,0001
  • Количество слоёв SSM: 24 слоя (в 2 раза больше, чем Transformer)
  • Размерность встраивания: 1024

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

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

1. Сравнение масштабируемости

• Предел Transformer: 2¹⁵ строк (ограничено 80GB видеопамятью)
• Предел Hydra: 2¹⁷ строк (ограничено 32-битным индексированием PyTorch, не аппаратным ограничением)
• Прирост производительности: Увеличение обрабатываемого размера данных в 100 раз

2. Сравнение предсказательной производительности

• Hydra vs Transformer: Средняя разница в точности -1,1%, разница в AUC -1,1%
• Hydra vs Mamba: Точность Hydra в среднем на 3,6% выше
• Анализ дисперсии: Hydra демонстрирует более низкую дисперсию производительности, чем Mamba

3. Анализ чувствительности к порядку

Измеренная дивергенцией KL:

• С увеличением количества перестановок RCP дивергенция KL значительно снижается
• Hydra демонстрирует более низкую чувствительность к порядку, чем Mamba
• Стратегия RCP эффективно снижает влияние аномальных перестановок

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

Влияние количества перестановок RCP

• Точность: Повышается с увеличением количества перестановок RCP, но величина улучшения относительно небольшая
• Дивергенция KL: Значительно снижается, указывая на уменьшение зависимости от порядка
• Вычислительные затраты: Линейно увеличиваются в r раз для времени вывода

Сравнение архитектур

• Однонаправленная vs двунаправленная: Двунаправленный механизм Hydra явно превосходит однонаправленную обработку Mamba
• Установка количества слоёв: Следуя рекомендациям статьи Mamba, используется в 2 раза больше слоёв, чем Transformer

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

1. Важность двунаправленности: Двунаправленная обработка критична для неупорядоченной природы табличных данных
2. Баланс эффективности и производительности: Достигнут значительный прирост эффективности при сохранении конкурентоспособной производительности
3. Эффективность RCP: Стратегия многократной перестановки и усреднения эффективно снижает чувствительность к порядку
4. Преодоление аппаратных ограничений: Успешно преодолены ограничения памяти Transformer при обработке крупномасштабных данных

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

Фундаментальные модели для табличных данных

• TabPFN: Пионерская модель Transformer для табличных данных
• TabFlex: Расширенный вариант с линейным вниманием
• Mambular: Модель глубокого обучения для табличных данных на основе Mamba

Модели пространства состояний

• Mamba: Селективная модель пространства состояний, обеспечивающая линейную сложность
• Hydra: Двунаправленное расширение SSM, поддерживающее некаузальное моделирование
• S4: Фундаментальная работа по структурированным моделям пространства состояний для последовательностей

Методы оптимизации эффективности

• FlashAttention: Снижение требований к памяти Transformer через оптимизацию ввода-вывода
• Linear Attention: Альтернативный механизм внимания с линейной сложностью

Выводы и обсуждение

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

1. Hydra успешно решает проблему масштабируемости TabPFN, повышая возможности обработки на два порядка
2. Двунаправленная SSM более подходит для неупорядоченной природы табличных данных, чем однонаправленная SSM
3. Стратегия RCP является эффективным методом снижения чувствительности SSM к порядку
4. Достигнута конкурентоспособная производительность с Transformer при сохранении линейной сложности

Ограничения

1. Требование переобучения: Из-за различий в архитектуре требуется переобучение всей модели
2. Ограничение контекста: Эксперименты по-прежнему ограничены 1000 строками, недостаточно исследованы крупномасштабные сценарии
3. Затраты RCP: Многократные перестановки увеличивают время вывода в r раз
4. Оптимизация порядка: Недостаточно глубоко исследованы оптимальные стратегии упорядочивания строк

Направления будущих исследований

1. Крупномасштабная валидация: Тестирование SSM-based TabPFN на наборах данных >10k строк
2. Оптимальное упорядочивание: Исследование оптимальных стратегий упорядочивания строк для SSM
3. Оптимизация архитектуры: Исследование более эффективных архитектур двунаправленной SSM
4. Теоретический анализ: Глубокое понимание теоретических основ двунаправленности для моделирования табличных данных

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

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

1. Чёткое определение проблемы: Точное выявление основного узкого места TabPFN и предложение целевого решения
2. Разумный выбор техники: Двунаправленная характеристика Hydra хорошо соответствует неупорядоченной природе табличных данных
3. Полный дизайн экспериментов: Включает многомерную оценку производительности, эффективности и чувствительности к порядку
4. Убедительные результаты: Достигнут значительный прирост масштабируемости при сохранении производительности
5. Высокая практическая применимость: Стратегия RCP проста, эффективна и легко реализуется и развёртывается

Недостатки

1. Ограниченная степень инноваций: В основном представляет собой комбинацию существующих технологий, не хватает фундаментальных инноваций
2. Недостаточный теоретический анализ: Отсутствует глубокое теоретическое объяснение того, почему двунаправленность решает проблему чувствительности к порядку
3. Ограниченный масштаб экспериментов: По-прежнему ограничены относительно небольшими наборами данных, недостаточно демонстрируют возможности крупномасштабной обработки
4. Неполное сравнение: Отсутствует прямое сравнение с другими методами линейной сложности (например, Linear Attention)
5. Недостаточный анализ гиперпараметров: Из-за высоких затрат на обучение не проведена достаточная оптимизация гиперпараметров

Влияние

1. Академический вклад: Предоставляет новые идеи и эмпирические доказательства для оптимизации эффективности фундаментальных моделей для табличных данных
2. Практическая ценность: Решает проблемы масштабируемости в практических приложениях, обладает высокой практической ценностью
3. Вдохновляющее значение: Демонстрирует потенциал SSM в моделировании структурированных данных, может вдохновить дальнейшие исследования
4. Воспроизводимость: Код открыто доступен, экспериментальная установка подробна, обладает хорошей воспроизводимостью

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

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

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

Основные цитируемые работы включают:

1. Hollmann et al. (2023) - исходная статья TabPFN
2. Gu & Dao (2023) - архитектура Mamba
3. Hwang et al. (2024) - двунаправленная SSM Hydra
4. Dao et al. (2022) - техника оптимизации FlashAttention
5. Zeng et al. (2024) - метод линейного внимания TabFlex

Данная статья вносит ценный вклад в решение проблемы масштабируемости фундаментальных моделей для табличных данных. Путём умелого сочетания двунаправленной SSM и стратегии повторяющихся перестановок успешно достигнут баланс между требованиями эффективности и производительности. Несмотря на некоторые недостатки в теоретических инновациях, её практическая ценность и вдохновляющее значение для будущих исследований заслуживают признания.