Введение в WaLa: Волновая Латентная Диффузия для 3D Генеративных Моделей

Тренировка генеративных моделей на больших объёмах 3D данных представляет собой значительный вызов. Сложность данных в 3D формате, обусловленная их кубической природой, приводит к резкому увеличению числа переменных, с которыми должна справляться модель, что значительно превышает сложность задач в области изображений или обработки естественного языка. К этому добавляются проблемы хранения и потоковой передачи данных, а также высокие затраты на использование облачных сервисов для обучения моделей на высокоразрешённых 3D датасетах. Кроме того, 3D формы могут быть представлены в различных форматах, таких как воксели, облака точек, сетки и неявные функции, каждый из которых имеет свои преимущества и недостатки в плане качества и компактности. Выбор оптимального представления для эффективного обучения и генерации остаётся открытым вопросом.

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

1. Использование низкоразрешённых представлений, таких как разреженные облака точек или грубые сетки, что уменьшает вычислительную сложность, но ограничивает способность модели точно захватывать детали.

2. Представление 3D форм через коллекцию 2D изображений или интеграцию изображений в обучающую функцию потерь, что позволяет модели учиться на внешнем виде объекта, но затрудняет захват внутренних деталей и увеличивает время обучения.

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

Волновые представления, такие как Neural Wavelet, UDiFF и волновые деревья, предлагают ряд преимуществ: данные могут быть легко сжаты с минимальной потерей деталей, а взаимосвязи между коэффициентами облегчают эффективное хранение, потоковую передачу и обработку больших 3D датасетов. Однако, несмотря на эти преимущества, волновые представления остаются довольно объёмными, особенно при масштабировании для крупных генеративных моделей.

В данной работе мы представляем WaLa (Wavelet Latent Diffusion) — новую структуру, которая дополнительно сжимает волновое представление для получения компактных латентных кодировок без значительной потери информации. Это позволяет эффективно масштабировать диффузионную генеративную модель в этом пространстве. Начав с усечённой функции знакового расстояния (TSDF) формы, мы сначала преобразуем её в 3D волновое дерево, затем обучаем VQ-VAE модель с адаптивной выборкой и сбалансированной настройкой для сжатия TSDF до компактной сетки, достигая впечатляющего соотношения сжатия 2427:1 при сохранении высокой точности реконструкции.

WaLa также поддерживает контролируемую генерацию через множество входных модальностей, делая фреймворк гибким и адаптируемым за пределы задач реконструкции 3D из одного вида. В результате наша модель генерирует 3D формы с сложной геометрией, правдоподобными структурами, сложными топологиями и гладкими поверхностями.

В этой статье мы рассмотрим концепции WaLa, его архитектуру, методы обучения, а также представим результаты и анализ производительности модели.