Flamingo Architecture › Architecture Overview

Flamingo consists of three main components arranged in a specific processing pipeline that turns raw pixels and text tokens into coherent multimodal representations:

1. Frozen Vision Encoder: A Normalizer-Free ResNet (NFNet) that extracts high-resolution visual features from input images. This encoder processes images through deep convolutional layers to produce rich spatial feature maps capturing visual semantics at multiple scales.

2. Perceiver Resampler: A trainable module that compresses variable-length vision features into a fixed number of visual tokens (typically 64). This compression is needed for computational tractability, letting the model to process high-resolution images without overwhelming the language model's attention mechanisms.

3. Frozen Language Model with GATED XATTN-DENSE: A pretrained Chinchilla language model augmented with trainable cross-attention layers inserted periodically throughout its depth. These inserted layers act as information channels, letting the language model to reference visual content without modifying its base parameters.

The architecture processes interleaved text and images by treating visual tokens as additional context that the language model can attend to through cross-attention, while the causal masking of the language model ensures autoregressive text generation remains intact. This design choice preserves the language model's ability to generate coherent, fluent text while grounding that generation in visual evidence.

The number of trainable parameters is surprisingly small relative to the total model size. In the largest Flamingo variant (80B parameters), the frozen vision encoder accounts for roughly 435M parameters, the frozen language model accounts for about 70B, and the trainable Perceiver Resampler plus GATED XATTN-DENSE layers account for only about 10B parameters. This means that roughly 87% of the model's parameters are never updated during multimodal training, which explains why the linguistic capabilities are so well preserved.

Vision encoders like NFNet output spatial feature maps representing thousands of local image regions. For a typical 224x224 image processed by an NFNet, the final convolutional layer might produce a 7x7 grid with 2,048-dimensional features, yielding 49 tokens. However, higher resolution images or different architectures can produce thousands of tokens. Feeding all of these into the language model at every layer would be computationally prohibitive, as the cross-attention mechanism would need to compute attention scores between every text token and every visual feature, resulting in memory requirements that scale with the product of text length and visual tokens.

The Perceiver Resampler solves this by learning a fixed number of latent queries that attend to the vision features, compressing the variable-length image representation into a constant number of visual tokens regardless of input resolution. This approach draws inspiration from the Perceiver architecture, which uses a similar mechanism to handle arbitrary input modalities by compressing them into a fixed-size latent bottleneck.

The mechanism works through three stages:

Learned Latents: The resampler maintains $N$ learned query vectors $Q \in \mathbb{R}^{N \times d}$, where $N$ is the desired number of output visual tokens (typically 64) and $d$ is the model dimension. Unlike positional encodings or static embeddings, these latents are optimized during training to extract the most salient visual information for language understanding. You can think of them as $N$ different "questions" the model asks about each image, where the questions themselves are learned rather than hand-designed.

Cross-Attention to Vision: These queries attend to the vision encoder outputs, which serve as keys $K$ and values $V$. If the vision encoder produces $M$ spatial features (where $M$ varies by image size), we have $K, V \in \mathbb{R}^{M \times d}$. The cross-attention operation effectively asks: given these $N$ learned queries, what information from the $M$ vision features is most relevant?

Self-Attention Among Latents: After the initial cross-attention, the resampler applies several layers of self-attention among the $N$ latent vectors, letting them to share information and produce a coherent compressed representation. This step ensures that the final visual tokens capture global image context rather than isolated local features, letting the tokens to collectively stand for the entire image content. Without this self-attention step, each latent would independently summarize the image from its own learned perspective, but the latents would not be aware of each other's "questions" or answers, potentially leading to redundancy or gaps in coverage.

Mathematically, for a single resampler layer, the output latents are computed by first applying cross-attention to the vision features, then applying self-attention among the latents:

$$\text{Latents}_{\text{out}} = \text{SelfAttn}(\text{CrossAttn}(\text{Latents}_{\text{in}}, \text{VisionFeatures}))$$

where:

• $\text{Latents}_{\text{in}}$: the learned query vectors $Q \in \mathbb{R}^{N \times d}$ before processing the current image
• $\text{VisionFeatures}$: the spatial features output by the vision encoder, serving as keys $K$ and values $V$
• $\text{CrossAttn}$: cross-attention operation where latents attend to vision features
• $\text{SelfAttn}$: self-attention operation among the $N$ latent vectors to share information
• $\text{Latents}_{\text{out}}$: the refined visual tokens after processing

This formulation compresses variable-length image features into a fixed number of tokens regardless of input resolution. The compression ratio can be substantial: for a vision transformer creating 1,024 tokens, compressing to 64 latents yields a 16-fold reduction in computation for downstream cross-attention operations.

The cross-attention operation follows the standard scaled dot-product formulation we covered in Part XIII, Ch 3: Scaled Dot-Product Attention:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

where:

• $Q$: the learned latent queries $\in \mathbb{R}^{N \times d}$ (the fixed set of queries attending to image features)
• $K$: keys derived from the vision encoder outputs $\in \mathbb{R}^{M \times d}$
• $V$: values derived from the vision encoder outputs $\in \mathbb{R}^{M \times d}$
• $d_k$: dimension of the keys (used for scaling to prevent softmax saturation)
• $\text{softmax}$: normalization function converting attention scores into a probability distribution over image locations

The scaling factor $\sqrt{d_k}$ prevents the dot products from becoming too large as dimensionality increases, which would push the softmax into regions with extremely small gradients. Without this scaling, the attention distribution would become extremely sharp, effectively attending to only one location and losing the soft, distributed attention that lets rich feature aggregation. The output is a fixed set of $N$ visual tokens that stand for the entire image, regardless of how many patches or regions the vision encoder initially produced.

One subtle but important aspect of the resampler is that the learned latent queries $Q$ are shared across all images processed by the model. This means the resampler applies the same "questions" to every image and learns to extract generally useful visual information rather than queries specialized to particular image categories. The power of the approach comes from the combination of these fixed learned questions with the flexible attention mechanism that routes information from different spatial locations depending on each specific image's content.

The resampler also handles the visual backbone's position encodings. The vision encoder may output features with 2D spatial structure (from convolutional layers) or with 1D position embeddings (from vision transformers). The resampler accepts either format, since it uses cross-attention rather than position-sensitive operations to aggregate information. This architectural flexibility allows Flamingo to use different vision backends without changing the downstream components.

While the Perceiver Resampler handles the vision-to-language dimensionality mismatch, we still need a mechanism to inject visual information into the language model. This must be done without destroying its pretrained linguistic capabilities. Simply adding visual tokens to the input sequence and using self-attention would work, but this approach would require modifying the language model's input embeddings and would not give the fine-grained control needed to prevent catastrophic forgetting. Flamingo adopts a more surgical approach: inserting cross-attention layers directly into the frozen language model at specific intervals.

These aren't standard cross-attention layers, however. They're wrapped in a advanced gating mechanism that ensures training stability and preserves the language model's initial behavior at the start of training. This gating acts as a safety valve, letting the model to gradually learn to use visual information rather than being forced to incorporate potentially noisy gradients from untrained cross-attention weights immediately.

Each GATED XATTN-DENSE block contains:

1. Layer Normalization applied to the input hidden states, stabilizing the distribution of activations before cross-modal interaction
2. Cross-Attention where text tokens (as queries) attend to visual tokens (as keys and values), creating the pathway for visual information to influence text generation
3. A gated residual connection with learnable scalar $\alpha$, controlling how much of the cross-attention output is added to the residual stream
4. Layer Normalization on the output of the cross-attention, preparing the mixed representations for further processing
5. A Dense (FFN) layer processing the mixed representations, letting nonlinear transformation of the visual-textual features
6. Another gated residual connection, independently controlling the contribution of the feed-forward transformation

The gating mechanism is important for preserving the frozen language model's behavior during initial training. Let $x$ be the input to the layer and $F(x)$ be the output of either the cross-attention or dense sublayer. The gated residual connection computes:

$$y = x + \tanh(\alpha) \cdot F(\text{LayerNorm}(x))$$

where:

• $x$: the input hidden states from the previous layer
• $F(x)$: the output of the cross-attention or dense sublayer after layer normalization
• $\alpha$: a learnable scalar parameter initialized to $0$
• $\tanh(\alpha)$: the gating factor that smoothly interpolates between identity ($0$) and full transformation ($\approx 1$)
• $y$: the output hidden states passed to the next layer

Since $\tanh(0) = 0$, at the beginning of training the layer computes $y = x$, effectively passing the input through unchanged. This initialization strategy acts as a warm start. This keeps the language model doesn't immediately forget its linguistic priors when visual inputs are introduced. As training progresses, $\alpha$ can grow (positively or negatively). This allows the layer to gradually learn to incorporate visual information without causing sudden distribution shifts in the frozen language model's activations.

The choice of $\tanh$ as the gating function deserves attention. A naive approach might use a linear gate, where $y = x + \alpha \cdot F(x)$. But a linear gate has no inherent bound, meaning $\alpha$ could grow unboundedly large during training, eventually overwhelming the residual stream and destroying the frozen model's behavior. Using $\tanh(\alpha)$ bounds the maximum gate value at 1.0. This keeps even with a very large $\alpha$, the cross-attention contribution never exceeds the magnitude of the identity path. This bounded behavior prevents training instability while still letting the model to learn strong visual conditioning when beneficial.

The GATED XATTN-DENSE also includes a separate feed-forward network with its own gating parameter. This FFN is not a standard transformer FFN; it doesn't process text features in isolation but rather processes the mixture of text and visual information that emerges from the cross-attention step. You can think of it as a learned transformation that converts the raw fusion of visual and textual features into a representation better aligned with the language model's expected input distribution. The independent gating of the FFN gives the model flexibility to learn that, for some tasks, simply injecting visual features through cross-attention is sufficient, while for others, a nonlinear transformation of the fused representation is necessary.

Flamingo doesn't add GATED XATTN-DENSE blocks after every transformer layer, as that would be computationally expensive and potentially redundant. Inserting cross-attention at every layer would increase the computational cost proportionally to the depth of the network, and empirical evidence suggests that visual information injected at early layers can propagate effectively through subsequent self-attention layers. Instead, it inserts them at regular intervals, typically every 4th or 6th layer in a 32-layer Chinchilla model. This sparse insertion pattern balances computational efficiency with the need for deep visual conditioning throughout the network. This keeps both low-level visual features and high-level semantic concepts can influence the language model's representations at appropriate levels of abstraction.

The specific choice of insertion frequency involves a tradeoff. Inserting cross-attention after every layer maximizes the model's ability to integrate visual information at fine-grained intervals but roughly doubles the number of cross-attention operations. Inserting every 8 layers reduces overhead materially but may limit the model's ability to tightly couple visual and textual reasoning across all levels of abstraction. The every-4th-layer configuration used in Flamingo was found empirically to strike a good balance.

When processing an interleaved sequence like:

The model processes text tokens autoregressively while maintaining cross-attention to the visual tokens extracted from each image. Importantly, each image's visual tokens attend only to that specific image, while text tokens can attend to all preceding images through the cross-attention mechanism. This architectural choice ensures that visual processing remains image-specific in the resampler, while the language model can freely associate and compare information across multiple images, letting capabilities like visual comparison, temporal reasoning, and multi-image context understanding.

The cross-attention masking implements a specific temporal causality rule: when generating a text token at position $t$, the GATED XATTN-DENSE layers allow attention only to images that appear before position $t$ in the interleaved sequence. This ensures that the model cannot "cheat" by attending to future images that haven't been described yet, maintaining the autoregressive generation property that makes the model useful in practice.