World Models (WMs) — это основная структура для разработки агентов, которые рассуждают и планируют в компактном скрытом пространстве. Однако обучение этих моделей непосредственно на основе пиксельных данных часто приводит к «коллапсу представлений», когда модель создаёт избыточные вложения, чтобы тривиально удовлетворить целям прогнозирования.
LeWorldModel (LeWM) — это первая JEPA (Joint-Embedding Predictive Architecture), которая стабильно обучается от начала до конца на основе необработанных пикселей, используя только два термина потерь:
* предсказание следующего вложения;
* регуляризатор, обеспечивающий гауссово распределение скрытых вложений.
Техническая архитектура и цель
LeWM состоит из двух основных компонентов, которые изучаются совместно:
* Encoder (zt=encθ (ot)) — преобразует необработанные пиксельные данные в компактное низкоразмерное скрытое представление. В реализации используется архитектура ViT-Tiny (~5M параметров).
* Predictor (Žt+1=predθ(zt, at)) — преобразователь (~10M параметров), который моделирует динамику среды, прогнозируя будущие скрытые состояния в зависимости от действий.
Модель оптимизируется с помощью упрощённой целевой функции, состоящей только из двух членов потерь:
$$\mathcal{L}{LeWM} \triangleq \mathcal{L}{pred} + \lambda SIGReg(Z)$$
Прогнозируемая потеря (Lpred) вычисляет среднеквадратичную ошибку (MSE) между прогнозируемыми и фактическими последовательными вложениями. SIGReg (Sketched-Isotropic-Gaussian Regularizer) — это термин, предотвращающий коллапс, который обеспечивает разнообразие признаков.
Согласно исследовательской статье, применение скорости выпадения (dropout rate) 0,1 в предикторе и специального шага проецирования (однослойный MLP с пакетной нормализацией) после энкодера критически важно для стабильности и производительности на последующих этапах.
Эффективность с помощью SIGReg и разреженной токенизации
Оценка нормальности в высокоразмерных скрытых пространствах — это серьёзная проблема масштабирования. LeWM решает эту задачу с помощью SIGReg, который использует теорему Крамера-Вольда: многомерное распределение соответствует цели (изотропному гауссовскому), если все его одномерные проекции соответствуют этой цели.
SIGReg проецирует скрытые вложения на M случайных направлений и применяет статистику теста Эппса-Пулли к каждой полученной одномерной проекции. Поскольку весовой коэффициент регуляризации λ является единственным эффективным гиперпараметром для настройки, исследователи могут оптимизировать его с помощью бинарного поиска со сложностью O(log n), что является значительным улучшением по сравнению с поиском за полиномиальное время (O(n6)), который требуется предыдущим моделям, таким как PLDM.
Бенчмарк скорости
В описанной настройке LeWM демонстрирует высокую вычислительную эффективность:
* Эффективность токенов: LeWM кодирует наблюдения примерно в 200 раз меньше токенов, чем DINO-WM.
* Скорость планирования: LeWM обеспечивает планирование до 48 раз быстрее, чем DINO-WM (0,98 с против 47 с за цикл планирования).
Свойства скрытого пространства и физическое понимание
Скрытое пространство LeWM поддерживает исследование физических величин и обнаружение физически неправдоподобных событий.
Нарушение ожиданий (VoE)
Используя фреймворк VoE, модель была оценена на её способность обнаруживать «сюрприз». Она присвоила более высокий уровень удивления физическим возмущениям, таким как телепортация; визуальные возмущения производили более слабые эффекты, а изменение цвета куба в OGBench-Cube не было значительным.
Выравнивание траектории
LeWM демонстрирует временное выравнивание скрытых траекторий, когда скрытые траектории естественным образом становятся более гладкими и линейными в процессе обучения. Примечательно, что LeWM достигает более высокой временной прямолинейности, чем PLDM, несмотря на отсутствие явного регуляризатора, поощряющего такое поведение.
| Характеристика | LeWorldModel (LeWM) | PLDM | DINO-WM | Dreamer / TD-MPC |
| — | — | — | — | — |
| Парадигма обучения | Стабильное обучение от начала до конца | Обучение от начала до конца | Замороженный базовый энкодер | Задача-специфическая |
| Тип входных данных | Необработанные пиксели | Необработанные пиксели | Пиксели (функции DINOv2) | Награды / привилегированное состояние |
| Термины потерь | 2 (прогнозирование + SIGReg) | 7 (на основе VICReg) | 1 (MSE по латентам) | Несколько (специфичных для задачи) |
| Настраиваемые гиперпараметры | 1 (эффективный вес λ) | 6 | N/A (фиксировано предварительной подготовкой) | Много (зависит от задачи) |
| Скорость планирования | До 48 раз быстрее | Быстро (компактные латенты) | Медленно (~50 раз медленнее, чем LeWM) | Варьируется (часто медленная генерация) |
| Предотвращение коллапса | Доказуемо (гауссовский априор) | Недостаточно определено / нестабильно | Ограничено предварительной подготовкой | Эвристический (например, реконструкция) |
| Требование | Задача-независимая / без вознаграждений | Задача-независимая / без вознаграждений | Замороженный предварительно обученный энкодер | Сигналы задачи / вознаграждения |
Ключевые выводы
1. Стабильное обучение от начала до конца. LeWM — это первая JEPA (Joint-Embedding Predictive Architecture), которая стабильно обучается от начала до конца на основе необработанных пикселей без необходимости использования эвристик, таких как stop-gradients, exponential moving averages (EMA) или замороженные предварительно обученные энкодеры.
2. Радикальная двухчленная цель. Процесс обучения упрощён до двух членов потерь — предсказания следующего вложения и регуляризатора SIGReg, — что сокращает количество настраиваемых гиперпараметров с шести до одного по сравнению с существующими альтернативами.
3. Построение для скорости в реальном времени. Представляя наблюдения примерно в 200 раз меньше токенов, чем у аналогов на основе фундаментальных моделей, LeWM планирует до 48 раз быстрее, выполняя полную оптимизацию траектории менее чем за одну секунду.
4. Доказуемое предотвращение коллапса. Чтобы модель не научилась создавать «мусорные» избыточные представления, в ней используется регуляризатор SIGReg; он использует теорему Крамера-Вольда, чтобы гарантировать, что высокоразмерные скрытые вложения остаются разнообразными и гауссово распределёнными.
5. Внутренняя физическая логика. Модель не просто прогнозирует данные; она улавливает значимую физическую структуру в своём скрытом пространстве, что позволяет ей точно исследовать физические величины и обнаруживать «невозможные» события, такие как телепортация объектов.
1. Какие основные проблемы решает модель LeWorldModel (LeWM) в контексте пиксельного моделирования мира?
В статье указано, что основная проблема, которую решает LeWM, — это коллапс представлений в моделях World Models (WMs) при обучении на основе пиксельных данных. Модель LeWM предотвращает коллапс представлений с помощью регуляризатора SIGReg, который обеспечивает гауссово распределение скрытых вложений.
2. Какие компоненты составляют архитектуру LeWM и какие функции они выполняют?
Архитектура LeWM состоит из двух основных компонентов: Encoder (преобразует необработанные пиксельные данные в компактное низкоразмерное скрытое представление) и Predictor (моделирует динамику среды, прогнозируя будущие скрытые состояния в зависимости от действий).
3. Какие преимущества предлагает LeWM по сравнению с другими моделями, такими как PLDM и DINO-WM?
Согласно статье, LeWM демонстрирует высокую вычислительную эффективность и скорость планирования. Модель кодирует наблюдения примерно в 200 раз меньше токенов, чем DINO-WM, и обеспечивает планирование до 48 раз быстрее, чем DINO-WM. Кроме того, LeWM стабильно обучается от начала до конца на основе необработанных пикселей без необходимости использования эвристик.
4. Какие методы используются в LeWM для предотвращения коллапса представлений?
В LeWM для предотвращения коллапса представлений используется регуляризатор SIGReg. Он использует теорему Крамера-Вольда для обеспечения гауссова распределения скрытых вложений и предотвращения создания моделью избыточных представлений.
5. Какие выводы можно сделать о потенциале LeWM для физического понимания и исследования физических величин?
Согласно статье, LeWM обладает потенциалом для физического понимания и исследования физических величин. Модель улавливает значимую физическую структуру в своём скрытом пространстве, что позволяет ей точно исследовать физические величины и обнаруживать «невозможные» события, такие как телепортация объектов.