Flamingo Architecture › Training Pipeline

Flamingo's training involves three distinct stages, each using different data sources but maintaining the same architectural principle: keep the vision encoder and language model frozen, train only the Perceiver Resampler and GATED XATTN-DENSE layers. This multi-stage curriculum progressively builds the model's capabilities, starting with general web-scale interleaved data, moving to higher-quality paired captions, and finally incorporating temporal video understanding.

The decision to use three separate stages rather than mixing all data jointly shows an important insight about how different data types contribute to different capabilities. Interleaved web data teaches contextual reasoning and few-shot pattern matching, since it contains natural demonstrations of how images relate to surrounding text in real documents. Paired caption data teaches fine-grained visual recognition, since high-quality captions tend to be more descriptive and accurate than casual textual references to images in web pages. Video data teaches temporal reasoning, since understanding dynamic scenes requires tracking how visual content changes across frames. Training on each type in sequence allows the model to build these skills progressively without the noisier data types overwhelming the signal from higher-quality sources.

The first stage trains on the Multimodal MassiveWeb dataset containing 183 million pages extracted from Common Crawl. Each training example consists of a sequence of up to 256 tokens of text interleaved with up to 5 images randomly sampled from the page. This stage teaches the model the basic skill of associating text with visual content in naturalistic contexts. The web data gives enormous diversity, covering everything from news articles and blog posts to instructional guides and forum discussions. This keeps the model learns reliable, generalizable associations between language and vision rather than domain-specific patterns.

The loss function is standard cross-entropy on the text tokens, measuring the negative log-likelihood of each token given the previous context and visual information:

$$\mathcal{L} = -\sum_{t} \log P(x_t \mid x_{<t}, \text{Images})$$

where:

• $\mathcal{L}$: the total cross-entropy loss across the sequence
• $x_t$: the target text token at position $t$
• $x_{<t}$: all preceding text tokens in the sequence (the autoregressive context)
• $\text{Images}$: all images in the sequence, accessed through cross-attention
• $P(x_t \mid x_{<t}, \text{Images})$: the probability assigned to the correct token by the model, conditioned on both text history and visual features

This objective teaches the model to predict text tokens while grounding predictions in visual content, forcing the cross-attention mechanisms to learn real associations between linguistic concepts and visual features. The model must learn to attend to relevant images when generating descriptive text, ignore irrelevant images when they don't support the current context, and integrate information across multiple images when the text references comparisons or sequences.

A subtle but important aspect of this training objective is that only text tokens contribute to the loss. Image tokens never appear as prediction targets. This means the model is never explicitly trained to describe images; it is instead trained to predict what text should surround images in natural documents. The visual grounding capabilities emerge indirectly, as a consequence of the model learning to make accurate text predictions that are consistent with visual context. This indirect training signal turns out to be remarkably effective, creating visual understanding that generalizes well to tasks like visual question answering even though the model was never explicitly trained on question-answer pairs.

The second stage incorporates 1.8 billion image-text pairs from the ALIGN dataset. While interleaved data teaches contextual reasoning and few-shot learning, paired data gives high-quality single-image captions that improve fine-grained visual recognition. ALIGN consists of image-alt_text pairs filtered from web pages using contrastive learning to ensure high quality and diversity.

During this stage, each example contains exactly one image and its corresponding caption, letting the model to focus on detailed visual understanding without the complexity of cross-image reasoning. This stage particularly improves the model's ability to recognize specific objects and attributes alongside spatial relationships, as the captions in ALIGN tend to be more descriptive and accurate than the surrounding text in arbitrary web documents.

The transition from interleaved to paired training is a shift in the complexity of the signal. Interleaved data contains implicit associations between images and text (the connection between an image and the nearby text must be inferred from context), while paired data contains explicit one-to-one correspondences. By first learning implicit associations and then reinforcing them with explicit pairings, the curriculum mirrors effective educational approaches: start with detailed contextual learning, then consolidate with targeted practice.

The final stage adds 27 million video-text pairs from the VideoTextPairs dataset. Videos are processed as sequences of frames (sampled at 1 FPS), with each frame treated as a separate image. This extends Flamingo's capabilities to temporal reasoning and video understanding, though it remains primarily an image-text model with video capabilities emerging from frame-level processing.

By treating video frames as individual images processed through the same Perceiver Resampler, Flamingo learns to track objects and actions across time, understand temporal transitions, and generate descriptions of dynamic scenes. This capability, while emergent rather than explicitly architected for video, shows the flexibility of the Flamingo framework in handling sequential visual data.

The video training stage also reveals an interesting limitation of the Perceiver Resampler design. Because each frame is independently processed through the resampler, the compressed visual tokens for each frame are computed independently, without any cross-frame attention. This means the resampler cannot directly learn motion-sensitive features that would require comparing adjacent frames. The temporal reasoning must emerge entirely through the language model's cross-attention to the sequence of per-frame visual tokens, which works reasonably well for slow or predictable motion but may miss fine-grained temporal details.

Training uses the AdamW optimizer with decoupled weight decay, similar to what we discussed in Part X, Ch 9: AdamW. The learning rate warms up linearly and then decays with a cosine schedule. Different components use different learning rates: the Perceiver Resampler and GATED XATTN-DENSE layers use a higher learning rate ($10^{-4}$ to $10^{-3}$) than the gating parameters ($10^{-5}$). This keeps stable initial training while letting the gates to open gradually.

This learning rate differentiation is important: if the gating parameters $\alpha$ were updated too quickly, the gates might swing open rapidly, causing sudden distribution shifts in the frozen language model's activations and potentially destabilizing training. By keeping the gating learning rate low, the model ensures smooth, gradual adaptation of the residual connections.

Gradient clipping is applied throughout training, as we discussed in Part X, Ch 13: Gradient Clipping. Because the frozen language model and vision encoder contribute gradients to the trainable layers through the chain rule (even though their own parameters aren't updated), the gradients flowing through the GATED XATTN-DENSE layers can be large and unstable early in training. Clipping to a maximum gradient norm of 1.0 prevents these large gradients from causing sudden parameter updates that could disrupt the gating mechanism.

The training uses a batch size that balances the three data sources. Specifically, the Flamingo paper reports using a batch that combines M3W, ALIGN, and VTP examples in ratios that give each data source roughly equal importance in the gradient updates despite their very different sizes. This weighting prevents the much larger ALIGN dataset from dominating training while still benefiting from its higher data quality.