A Statistical Theory of Contrastive Pre-training and Multimodal Generative AI
Kazusato Oko, Licong Lin, Yuhang Cai, Song Mei
TL;DR
This work addresses why contrastive pre-training yields robust, transferable representations for multimodal AI. It introduces the concept of approximate sufficient statistics and a Joint Generative Hierarchical Model (JGHM), showing near-minimizers of the CLIP loss implicitly capture sufficient statistics and enabling end-to-end sample-efficient learning with transformers. The authors prove excess-risk bounds for downstream tasks like zero-shot classification and conditional diffusion models, and demonstrate that transformers can efficiently approximate belief propagation within GHMs. Numerical simulations corroborate the theory, illustrating strong generalization and adaptability of pre-trained CLIP representations to ZSC, CDMs, and VLMs, while highlighting the practical impact on sample efficiency and robustness in multimodal generation and understanding.
Abstract
Multi-modal generative AI systems, such as those combining vision and language, rely on contrastive pre-training to learn representations across different modalities. While their practical benefits are widely acknowledged, a rigorous theoretical understanding of the contrastive pre-training framework remains limited. This paper develops a theoretical framework to explain the success of contrastive pre-training in downstream tasks, such as zero-shot classification, conditional diffusion models, and vision-language models. We introduce the concept of approximate sufficient statistics, a generalization of the classical sufficient statistics, and show that near-minimizers of the contrastive pre-training loss are approximately sufficient, making them adaptable to diverse downstream tasks. We further propose the Joint Generative Hierarchical Model for the joint distribution of images and text, showing that transformers can efficiently approximate relevant functions within this model via belief propagation. Building on this framework, we derive sample complexity guarantees for multi-modal learning based on contrastive pre-trained representations. Numerical simulations validate these theoretical findings, demonstrating the strong generalization performance of contrastively pre-trained transformers in various multi-modal tasks.