Оценка распределения данных — фундаментальная задача для генеративных моделей и байес-вывода. Но до сих пор инженерам приходилось выбирать: либо универсальные, но неточные методы, либо точные, но требующие переобучения под каждое новое распределение. Представлен DiScoFormer (Density and Score Transformer) — трансформер, который решает эту дилемму.
DiScoFormer оценивает плотность (smooth histogram, показывает, где точки данных сгруппированы) и скоринг (градиент логарифма плотности, указывает направление, в котором плотность растёт быстрее всего) распределения данных за один проход, не требуя переобучения под каждое новое распределение. Эта универсальность важна, поскольку скоринг используется в основе диффузионных моделей (как Stable Diffusion), байесовского семплирования и симуляций частиц.
Существующие подходы имеют свои ограничения:
- Kernel Density Estimation (KDE) — не требует обучения и применима к любым распределениям, но её точность резко падает с ростом размерности данных.
- Нейросетевые модели скоринга — точны в высоких размерностях, но каждую такую модель нужно обучать с нуля для каждого нового распределения.
DiScoFormer устраняет этот компромисс. Модель использует слои трансформерных блоков и кросс-внимание (cross-attention), что позволяет ей оценивать плотность и скоринг в любой точке, а не только там, где есть данные. Плотность и скоринг математически связаны (скоринг — это градиент логарифма плотности), и DiScoFormer использует это: у неё общий бэкенд с двумя выходными головами — для плотности и для скоринга. Такая связь также позволяет использовать потери согласованности (consistency loss) без меток, адаптируя модель к данным вне распределения без необходимости в новых обучающих данных.
Архитектура трансформера хорошо подходит для этой задачи. Аналитически показано, что веса одного блока внимания почти совпадают с гауссовым ядром KDE, что делает трансформер функциональным обобщением ядерных методов. Модель включает KDE как частный случай и улучшает его, обучаясь нескольким масштабам одновременно и адаптируя их к данным. Для обучения DiScoFormer использовались гауссовы смешанные модели (GMMs), которые являются универсальными аппроксиматорами плотности.
DiScoFormer предлагает универсальное решение для оценки распределений, которое может упростить работу с генеративными моделями и байесовским выводом. Модель устраняет компромисс между универсальностью и необходимостью переобучения, но её широкое внедрение и производительность в различных реальных сценариях ещё предстоит оценить.