DiScoFormer: один трансформер для плотности и сметки, на любых распределениях
Много задач машинного обучения и естественных наук требуют одного: по набору данных восстановить распределение, из которого они пришли. Нужны две величины: плотность (как гистограмма, но гладкая) и сметка, градиент логарифма плотности, указывающий направление, куда плотность растёт быстрее. Диффузионные модели (они лежат в основе Stable Diffusion и DALL-E) ходят по сметке, превращая шум в реалистичное изображение.
Сейчас нужно выбирать: классическая оценка плотности ядра (KDE) работает на любых распределениях без обучения, но теряет точность в высоких размерностях; нейросетевые методы точны в высоких размерностях, но требуют переобучения на каждое новое распределение.
Allen Institute представила DiScoFormer (Density and Score Transformer), один трансформер, который за одну подачу вперёд оценивает обе величины без переобучения. Архитектура опирается на кросс-внимание и использует математическую связь между плотностью и сметкой: они делят общий хребет, но выводят две головки. На этапе вывода модель адаптируется к новому распределению в полёте, улучшая себя через несколько шагов градиентного спуска на согласованность.
Математически трансформер обобщает KDE: одна голова внимания, это почти ядро Гаусса, поэтому блок кросс-внимания уже воспроизводит KDE, но дальше модель учится несколько таких масштабов одновременно и адаптирует их к данным.
Для обучения исследователи использовали смеси Гауссов, они универсальные аппроксиматоры плотности и имеют точные формулы для плотности и сметки. На каждой итерации генерируется новая смесь, давая модели фактически бесконечные примеры с известными ответами.
На практике DiScoFormer везде бьёт KDE. В 100 размерностях разница огромная: ошибка сметки меньше в 6.5 раз, ошибка плотности в 37 раз, и модель продолжает улучшаться с добавлением примеров, тогда как KDE исчерпывает память. DiScoFormer обобщается на смеси с большим числом мод, чем видел на обучении, и на негауссовские формы вроде Лапласа и t-распределения. Единственное преимущество KDE, скорость на малых данных.
Исследователи видят перспективу: оценка сметки, ключевая задача в диффузионном моделировании, байесовском выводе и вычислительных науках. Претренированный, повторно используемый оценщик, точный в высоких размерностях, мог бы сократить обучение во всех этих областях сразу.
Ключевые факты
- DiScoFormer оценивает плотность и градиент лог-плотности за один проход, без переобучения на новые распределения
- В 100-мерном пространстве превосходит KDE в 6.5, 37 раз, остаётся точной в высоких размерностях
- Архитектура обобщает классическую оценку плотности ядра: кросс-внимание, это параметризованное ядро Гаусса
- Использует математическую связь между плотностью и сметкой для адаптации на новых распределениях без переобучения
- Применяется в диффузионных моделях, байесовском выводе, научных симуляциях, везде, где нужна оценка распределения
Почему это важно
Оценка распределения, базовая задача в ML и естественных науках. Сметка (градиент логарифма плотности), это то, по чему ходит диффузия, байесовский вывод, плазменные симуляции. До сих пор приходилось выбирать: быстрый, но неточный KDE или точный, но медленный метод для каждого нового распределения. DiScoFormer ломает этот компромисс, один трансформер работает везде, точен в высоких размерностях и не требует переобучения.
Кому это важно
Исследователям в ML, байесовском выводе, физике и вычислительных науках, которые работают с оценкой распределений. Разработчикам диффузионных моделей и генеративных систем. Всем, кто сейчас переобучает свои модели на каждое новое распределение или вынужден жертвовать точностью ради универсальности.
Как это применить
DiScoFormer, это компонент, который встраивается в pipeline. Даёшь ему набор данных, на выходе получаешь плотность и сметку. На практике это замена KDE в местах, где нужна точность. Модель адаптируется в полёте через несколько градиентных шагов на согласованность между плотностью и её градиентом, можно даже улучшать оценку на нераспределённых примерах без переобучения.
Можно ли доверять
Работа из Allen Institute, опубликована как техреспорт. Математика опирается на связь между плотностью и сметкой (дифференцирование). Обучение на смесях Гауссов, стандартный приём: смеси, универсальные аппроксиматоры, есть точные формулы. Тесты на Гауссовых смесях, Лапласе, t-распределениях показывают, что моделирования выполнены честно. Адаптация в полёте через градиентные шаги, верная идея, аналогичная известным техникам. Результаты убедительны.
Риски и подводные камни
KDE по-прежнему быстрее на малых наборах данных, поэтому замена не универсальна. DiScoFormer требует трансформера, инфраструктуры, которой нет везде. Обучение на смесях Гауссов означает, что точность может падать на очень экзотичных распределениях, не похожих на смеси. Адаптация в полёте требует доступа к градиентам, не всегда возможно в production. Настройка гиперпараметров (число шагов адаптации, размер модели), открытый вопрос.
«Претренированный, plug-in оценщик, который остаётся точным в высоких размерностях и исключает необходимость переобучения для каждой задачи, мог бы сократить затраты сразу во всех этих областях, один модель, переиспользуемая везде, где появляются сметка и плотность.»
— Allen Institute (техреспорт DiScoFormer, Hugging Face Blog)