2025-11-25T22:34:18.624435

Efficient Autoregressive Inference for Transformer Probabilistic Models

Hassan, Loka, Li et al.
Transformer-based models for amortized probabilistic inference, such as neural processes, prior-fitted networks, and tabular foundation models, excel at single-pass marginal prediction. However, many real-world applications, from signal interpolation to multi-column tabular predictions, require coherent joint distributions that capture dependencies between predictions. While purely autoregressive architectures efficiently generate such distributions, they sacrifice the flexible set-conditioning that makes these models powerful for meta-learning. Conversely, the standard approach to obtain joint distributions from set-based models requires expensive re-encoding of the entire augmented conditioning set at each autoregressive step. We introduce a causal autoregressive buffer that preserves the advantages of both paradigms. Our approach decouples context encoding from updating the conditioning set. The model processes the context once and caches it. A dynamic buffer then captures target dependencies: as targets are incorporated, they enter the buffer and attend to both the cached context and previously buffered targets. This enables efficient batched autoregressive generation and one-pass joint log-likelihood evaluation. A unified training strategy allows seamless integration of set-based and autoregressive modes at minimal additional cost. Across synthetic functions, EEG signals, cognitive models, and tabular data, our method matches predictive accuracy of strong baselines while delivering up to 20 times faster joint sampling. Our approach combines the efficiency of autoregressive generative models with the representational power of set-based conditioning, making joint prediction practical for transformer-based probabilistic models.
academic

Эффективный Авторегрессивный Вывод для Вероятностных Моделей на основе Трансформеров

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

  • ID статьи: 2510.09477
  • Название: Efficient Autoregressive Inference for Transformer Probabilistic Models
  • Авторы: Conor Hassan, Nasrulloh Loka, Cen-You Li, Daolang Huang, Paul E. Chang, Yang Yang, Francesco Silvestrin, Samuel Kaski, Luigi Acerbi
  • Классификация: stat.ML cs.LG
  • Дата публикации: 10 октября 2025 г. (препринт arXiv)
  • Ссылка на статью: https://arxiv.org/abs/2510.09477

Аннотация

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

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

Основная проблема

Существующие вероятностные модели на основе трансформеров сталкиваются с фундаментальным узким местом эффективности: при необходимости генерации совместного распределения требуется переодирование всего условного множества на каждом авторегрессивном шаге. Конкретно:

  1. Ограничения моделей с условием на множество: Нейронные процессы (NPs), сети с предварительной подгонкой (PFNs) и подобные модели хорошо работают при маргинальном предсказании, но при авторегрессивном развертывании требуют повторного кодирования контекста, что приводит к вычислительной сложности O(K(N+K)²)
  2. Недостатки чистых авторегрессивных моделей: Хотя они вычислительно эффективны, им не хватает гибкой способности условия на множество, что ограничивает их применение в задачах метаобучения

Значимость

Предсказание совместного распределения критически важно для нескольких ключевых приложений:

  • Временные зависимости при интерполяции сигналов
  • Корреляции признаков при предсказании нескольких столбцов
  • Последовательные зависимости при моделировании поведенческих данных
  • Оценка совместного правдоподобия при байесовском выборе модели

Ограничения существующих методов

  1. Авторегрессивное развертывание TNP-D: На каждом шаге требуется переодирование растущего условного множества
  2. TNP-A: Требует обработки повторяющихся целевых множеств как при обучении, так и при выводе, что приводит к огромным вычислительным затратам
  3. TNP-ND: Ограничено многомерными гауссовыми распределениями, имеет ограниченную выразительную способность

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

  1. Предложен механизм причинного авторегрессивного буфера: Разделяет кодирование контекста с условием на множество и последовательное предсказание, обеспечивая эффективную совместную выборку и оценку правдоподобия
  2. Разработана единая стратегия обучения: Использует маскированное внимание и обучение с расписанием размера буфера, позволяя единой модели с минимальными дополнительными затратами изучить оба режима работы
  3. Проверена широкая применимость: Достигнуто ускорение совместной выборки до 20 раз на TNPs/PFNs и табличных фундаментальных моделях при сохранении сравнимой точности предсказания
  4. Оптимизирована теоретическая сложность: Вычислительная сложность снижена с O(K(N+K)²) до O(N²+NK+K²)

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

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

Дано контекстное множество C = {(xₙ, yₙ)}ᴺₙ₌₁ и целевое множество T = {(xₘ, yₘ)}ᴹₘ₌₁, цель состоит в изучении предсказывающего распределения p_θ(y₁:ₘ|x₁:ₘ; C), где θ — параметры модели.

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

Основные компоненты

  1. Кодировщик контекста rC: Обрабатывает пары контекста, использует двусторонний многоголовый самовнимание, кэширует пары ключ-значение на каждом слое
  2. Кодировщик буфера rB: Применяет строго причинное многоголовое самовнимание к префиксу буфера
  3. Декодер целевых переменных rtgt: Запрашивает кэшированный контекст и видимый префикс буфера через перекрестное внимание

Параметризация предсказывающего распределения

p_θ(y*₁:K|x*₁:K; C) = ∏ᴷₖ₌₁ p_θ(y*ₖ|rtgt(x*ₖ, [rC(C), b₁:ₖ₋₁]))

где bₖ = rB((xₖ, yₖ), rC(C), b₁:ₖ₋₁)

Проектирование маски внимания

Реализует четыре ключевых требования:

  • (R1) Неизменяемость контекста: кодируется один раз и кэшируется как только для чтения
  • (R2) Строгая причинность буфера: токен j может внимать только позициям <j
  • (R3) Однонаправленный поток информации из контекста: нет обратной записи в C
  • (R4) Целевые переменные внимают кэшированному контексту и видимому префиксу буфера

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

1. Разделенная конструкция

  • Статический кэш контекста: Кодируется один раз, переиспользуется многократно
  • Динамический буфер: Инкрементальное обновление, захватывает зависимости между целевыми переменными

2. Расписание обучения

  • 50% целевых переменных внимают только контексту
  • 50% целевых переменных внимают контексту + префиксу буфера случайной длины
  • Гарантирует хорошую работу модели при различных состояниях буфера

3. Эффективные режимы вывода

  • Авторегрессивная выборка: Предварительное заполнение контекста, последовательное декодирование целевых переменных
  • Оценка совместного правдоподобия: Единый прямой проход вычисляет все условные вероятности
  • Пакетная выборка: Совместный кэш контекста, независимые состояния буфера

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

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

  1. Синтетические функции:
    • Гауссовские процессы (GP): ядра RBF, Matérn-3/2, Matérn-5/2
    • Пилообразные функции: негауссовы, разрывные производные
  2. Данные ЭЭГ: 11 520 испытаний, 122 субъекта, 7 релевантных каналов, 256 временных точек
  3. Модель мультисенсорного причинного вывода: Данные аудиовизуального локализационного эксперимента, 15 участников
  4. Табличные данные: Наборы данных UCI (потребление электроэнергии, выбросы газотурбинных установок, совместное использование велосипедов)

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

  • Среднее логарифмическое правдоподобие: Оценка качества предсказания
  • Время на стене: Фактическое время выполнения выборки, оценки правдоподобия, шагов обучения
  • RMSE логарифмического маргинального правдоподобия: Точность при выборе модели

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

  • TNP-D-Ind: Независимые предсказания, быстро, но без моделирования зависимостей
  • TNP-D-AR: Авторегрессивное развертывание, выразительно, но требует переодирования
  • TNP-ND: Совместное многомерное гауссово распределение, ограниченная выразительность
  • TNP-A: Полностью авторегрессивное моделирование, медленно при обучении и выборке

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

  • Оптимизатор: Adam, скорость обучения 1×10⁻⁴
  • Архитектура: 6-слойный трансформер, 4 головы внимания, размерность 128
  • Головка предсказания: Гауссова смесь из 20 компонент
  • Размер буфера: K=16 (основные эксперименты)

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

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

Вычислительная эффективность

  • Авторегрессивная выборка: В 3-20 раз быстрее, чем TNP-A и TNP-D-AR
  • Оценка правдоподобия: Сравнима с TNP-A, в K раз быстрее, чем TNP-D-AR
  • Скорость обучения: В 4-12 раз быстрее, чем TNP-A, сравнима с самыми быстрыми базовыми методами

Точность предсказания

Набор данныхTNP-D-ARTNP-AДанный метод (K=16)Данный метод (K=1)
GP2.570.802.512.56
Пилообразная1.05-0.431.001.09
ЭЭГ-Инт0.510.460.520.54
ЭЭГ-Пред1.07-0.040.851.21

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

  • Влияние размера буфера: K=1 эквивалентен стандартной авторегрессии, K=16 показывает небольшое снижение производительности, но значительное ускорение
  • Пользовательские ядра Triton: Обеспечивают значительное ускорение при больших пакетах
  • Режимы внимания: Даже при отключении FlashAttention TNP-A остается на несколько порядков медленнее других методов

Анализ конкретных случаев

В задаче мультисенсорного причинного вывода:

  • Выбор модели: RMSE логарифмического маргинального правдоподобия 3.56, близко к 3.47 для TNP-D-AR
  • Предсказание данных: Среднее логарифмическое правдоподобие -2.76, сравнимо со всеми сильными базовыми методами
  • Корреляция с истинными значениями: R²=1.00 (логарифмическое маргинальное правдоподобие), R²=0.92 (ΔLML)

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

Нейронные процессы и сети с предварительной подгонкой

Данный метод служит модульным компонентом, который может быть интегрирован в существующие архитектуры NP/PFN. Дополняет предыдущие работы, сосредоточенные на масштабируемости контекстного множества, сосредотачиваясь на эффективности авторегрессивной совместной выборки.

Вероятностные модели на основе трансформеров

Строится на тренде представления байесовского вывода как задач контекстного обучения, используя варианты NP и PFN на основе трансформеров.

Табличные фундаментальные модели

Естественно интегрируется с моделями типа TabPFN и TabICL, обеспечивая дополнительный модуль для эффективного совместного предсказания.

Авторегрессивная оценка совместной плотности

Связано с TNP-A, но имеет ключевые отличия: TNP-A использует повторение целевых переменных как при обучении, так и при выводе, тогда как данный метод требует это только при оценке правдоподобия.

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

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

  1. Прорыв в эффективности: Успешно внедрена эффективность авторегрессивных трансформеров в рамки NP/PFN
  2. Сохранение производительности: Значительное ускорение при сохранении точности предсказания
  3. Широкая применимость: Метод проверен на эффективность в нескольких областях и задачах

Ограничения

  1. Масштабирование длины буфера: При увеличении K остается O(K²) член, текущая реализация использует фиксированные позиционные кодирования
  2. Дрейф качества при длинных буферах: Может быть снижение качества по сравнению с точной авторегрессией с переодированием на каждом шаге
  3. Использование памяти: Требует поддержания кэша контекста и состояния буфера

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

  1. Улучшение позиционного кодирования: Использование RoPE или ALiBi для поддержки более длинных последовательностей
  2. Спекулятивное декодирование: Адаптивные стратегии вывода, заимствованные из процессов draft-verify
  3. Параметрически эффективная тонкая настройка: Использование адаптеров или LoRA для добавления функциональности буфера к предварительно обученным моделям

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

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

  1. Высокая инновационность: Умно решает компромисс между условием на множество и авторегрессивной эффективностью
  2. Прочная теория: Обеспечивает четкий анализ сложности и математические выводы
  3. Комплексные эксперименты: Охватывает синтетические данные, реальные данные, несколько областей приложений
  4. Инженерная оптимизация: Включает пользовательские CUDA ядра и другие низкоуровневые оптимизации
  5. Воспроизводимость: Предоставляет подробные детали реализации и планирует открыть исходный код

Недостатки

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

Влияние

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

Сценарии применения

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

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

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

  • Garnelo et al. (2018): Оригинальная статья Neural Processes
  • Nguyen & Grover (2022): Transformer Neural Processes
  • Müller et al. (2022): Prior-Fitted Networks
  • Bruinsma et al. (2023): Autoregressive Conditional Neural Processes
  • Jingang et al. (2025): Табличная фундаментальная модель TabICL

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