MoMa: Эффективное предварительное обучение с ранним объединением смеси экспертов, осведомленных о модальности

MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts

July 31, 2024

Авторы: Xi Victoria Lin, Akshat Shrivastava, Liang Luo, Srinivasan Iyer, Mike Lewis, Gargi Gosh, Luke Zettlemoyer, Armen Aghajanyan

cs.AI

Аннотация

Мы представляем MoMa, новую архитектуру модальности-осведомленного смешанного экспертного (MoE) подхода, разработанную для предварительного обучения смешанных, раннего слияния языковых моделей. MoMa обрабатывает изображения и текст в произвольной последовательности, разделяя экспертные модули на модально-специфические группы. Эти группы исключительно обрабатывают назначенные токены, используя изученное маршрутизирование внутри каждой группы для поддержания семантически информированной адаптивности. Наши эмпирические результаты показывают значительное увеличение эффективности предварительного обучения благодаря этому модально-специфическому распределению параметров. При бюджете обучения в 1 трлн токенов модель MoMa 1.4B, включающая 4 текстовых эксперта и 4 эксперта по изображениям, достигает впечатляющих сбережений FLOPs: общий показатель уменьшается в 3.7 раза, сокращение для текста составляет 2.6 раза, а для обработки изображений - 5.2 раза по сравнению с плотной базовой моделью с эквивалентным вычислительным объемом, измеренным по потере предварительного обучения. Это превосходит стандартный выбор эксперта MoE с 8 смешанными модальными экспертами, который достигает сбережения FLOPs в 3 раза (3 для текста, 2.8 для изображения). Комбинирование MoMa с смесью глубин (MoD) дополнительно улучшает сбережения FLOPs предварительного обучения до 4.2 раза в общем (текст: 3.4 раза, изображение: 5.3 раза), хотя эта комбинация вредит производительности в причинном выводе из-за увеличенной чувствительности к точности маршрутизатора. Эти результаты демонстрируют потенциал MoMa значительно улучшить эффективность предварительного обучения смешанных, раннего слияния языковых моделей, открывая путь к более ресурсоэффективным и способным мультимодальным системам искусственного интеллекта.

English

We introduce MoMa, a novel modality-aware mixture-of-experts (MoE) architecture designed for pre-training mixed-modal, early-fusion language models. MoMa processes images and text in arbitrary sequences by dividing expert modules into modality-specific groups. These groups exclusively process designated tokens while employing learned routing within each group to maintain semantically informed adaptivity. Our empirical results reveal substantial pre-training efficiency gains through this modality-specific parameter allocation. Under a 1-trillion-token training budget, the MoMa 1.4B model, featuring 4 text experts and 4 image experts, achieves impressive FLOPs savings: 3.7x overall, with 2.6x for text and 5.2x for image processing compared to a compute-equivalent dense baseline, measured by pre-training loss. This outperforms the standard expert-choice MoE with 8 mixed-modal experts, which achieves 3x overall FLOPs savings (3x for text, 2.8x for image). Combining MoMa with mixture-of-depths (MoD) further improves pre-training FLOPs savings to 4.2x overall (text: 3.4x, image: 5.3x), although this combination hurts performance in causal inference due to increased sensitivity to router accuracy. These results demonstrate MoMa's potential to significantly advance the efficiency of mixed-modal, early-fusion language model pre-training, paving the way for more resource-efficient and capable multimodal AI systems.