Flamingo Architecture › Implementation Walkthrough
Language AI HandbookAll chapters Flamingo Architecture › Implementation WalkthroughOne section · Sections in this chapter6 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement the core components of the Flamingo architecture. We'll construct a simplified version that shows the Perceiver Resampler and Gated Cross-Attention mechanisms, then show how they integrate with a pretrained language model.
import torch
# Set seed for reproducibility
torch.manual_seed(42)Perceiver Resampler Implementation
The Perceiver Resampler uses learned latent queries to compress vision features. We'll implement it with configurable depth (number of layers) and the ability to handle variable input lengths.
class PerceiverResampler(nn.Module):
"""
Compresses variable-length vision features to fixed number of tokens.
Args:
dim: Model dimension (d_model)
num_latents: Number of output visual tokens (N)
num_layers: Number of resampler layers
num_heads: Number of attention heads
dim_head: Dimension per head
"""
def __init__(
self,
dim: int = 512,
num_latents: int = 64,
num_layers: int = 6,
num_heads: int = 8,
dim_head: int = 64,
):
super().__init__()
self.num_latents = num_latents
self.dim = dim
# Learned latent queries: N x D
self.latents = nn.Parameter(torch.randn(num_latents, dim) * 0.02)
# Project vision features to model dimension if needed
# (Assuming vision encoder outputs dim-dimensional features)
self.to_q = nn.Linear(dim, dim, bias=False)
self.to_kv = nn.Linear(dim, dim * 2, bias=False)
self.to_out = nn.Linear(dim, dim)
# Self-attention among latents
self.self_attn_norm = nn.LayerNorm(dim)
self.self_attn = nn.MultiheadAttention(
embed_dim=dim, num_heads=num_heads, batch_first=True
)
# Feed-forward network for latents
self.ffn_norm = nn.LayerNorm(dim)
self.ffn = nn.Sequential(
nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim)
)
self.num_layers = num_layers
def forward(self, vision_features: torch.Tensor) -> torch.Tensor:
"""
Args:
vision_features: (batch, num_vision_tokens, dim)
Returns:
visual_tokens: (batch, num_latents, dim)
"""
batch_size = vision_features.shape[0]
# Expand latents to batch size: (batch, N, D)
latents = self.latents.unsqueeze(0).expand(batch_size, -1, -1)
# Cross-attention: Latents attend to vision features
q = self.to_q(latents)
k, v = self.to_kv(vision_features).chunk(2, dim=-1)
# Scaled dot-product attention
scale = math.sqrt(self.dim)
attn_scores = torch.matmul(q, k.transpose(-2, -1)) / scale
attn_probs = F.softmax(attn_scores, dim=-1)
cross_attn_out = torch.matmul(attn_probs, v)
cross_attn_out = self.to_out(cross_attn_out)
# Residual connection for cross-attention
latents = latents + cross_attn_out
# Self-attention layers among latents
for _ in range(self.num_layers):
# Self-attention
normed = self.self_attn_norm(latents)
attn_out, _ = self.self_attn(normed, normed, normed)
latents = latents + attn_out
# FFN
normed = self.ffn_norm(latents)
latents = latents + self.ffn(normed)
return latentsInput vision shape: torch.Size([2, 196, 512]) Output visual tokens shape: torch.Size([2, 64, 512]) Compression ratio: 3.1x
The resampler successfully compresses 196 vision tokens into 64 visual tokens, a compression ratio of roughly 3x. This compression is important for managing computational costs when processing multiple high-resolution images. Without this compression, processing a sequence of 10 images would require the language model to attend to 1,960 visual tokens, creating a significant computational bottleneck in the cross-attention layers. The 9.4x reduction in attention cost comes from the quadratic nature of attention: halving the number of tokens reduces attention cost by a factor of four, and reducing by 3.1x reduces it by more than nine.
Gated Cross-Attention Layer
Now let's implement the GATED XATTN component. This layer allows text tokens to attend to visual tokens while using the gating mechanism to control information flow.
class GatedCrossAttention(nn.Module):
"""
Cross-attention from text (queries) to vision (keys/values) with gating.
"""
def __init__(
self,
dim: int = 512,
num_heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
):
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.scale = math.sqrt(dim_head)
self.norm = nn.LayerNorm(dim)
self.to_q = nn.Linear(dim, dim, bias=False)
self.to_kv = nn.Linear(dim, dim * 2, bias=False)
self.to_out = nn.Sequential(nn.Linear(dim, dim), nn.Dropout(dropout))
# Learnable gating parameter, initialized to 0
# tanh(0) = 0, so initially gate is closed (no visual influence)
self.alpha = nn.Parameter(torch.tensor(0.0))
def forward(
self,
text_hidden: torch.Tensor,
visual_tokens: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Args:
text_hidden: (batch, seq_len, dim) - text representations
visual_tokens: (batch, num_visual, dim) - visual tokens from resampler
attention_mask: Optional mask for padding
Returns:
gated_output: (batch, seq_len, dim)
"""
batch_size, seq_len, _ = text_hidden.shape
num_visual = visual_tokens.shape[1]
# Pre-normalization
normed = self.norm(text_hidden)
# Project to Q, K, V
q = self.to_q(normed)
k, v = self.to_kv(visual_tokens).chunk(2, dim=-1)
# Reshape for multi-head attention: (batch, heads, seq, dim_head)
head_dim = self.dim // self.num_heads
q = q.view(batch_size, seq_len, self.num_heads, head_dim).transpose(
1, 2
)
k = k.view(batch_size, num_visual, self.num_heads, head_dim).transpose(
1, 2
)
v = v.view(batch_size, num_visual, self.num_heads, head_dim).transpose(
1, 2
)
# Scaled dot-product attention
attn_scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale
if attention_mask is not None:
# Expand mask for heads: (batch, 1, 1, num_visual)
mask = attention_mask.unsqueeze(1).unsqueeze(2)
attn_scores = attn_scores.masked_fill(mask == 0, float("-inf"))
attn_probs = F.softmax(attn_scores, dim=-1)
attn_out = torch.matmul(attn_probs, v)
# Reshape back: (batch, seq_len, dim)
attn_out = (
attn_out.transpose(1, 2)
.contiguous()
.view(batch_size, seq_len, self.dim)
)
attn_out = self.to_out(attn_out)
# Gated residual connection: y = x + tanh(alpha) * F(x)
gate = torch.tanh(self.alpha)
output = text_hidden + gate * attn_out
return outputInitial alpha: 0.0000 Initial gate (tanh(alpha)): 0.0000 Output shape: torch.Size([2, 32, 512]) Output equals input (initially): True After training, gate: 0.9640
Initially, the gate is effectively closed (), meaning the output equals the input. This preserves the language model's behavior at the start of training. As increases (simulated here by setting it to 2.0), the gate opens to approximately 0.96, letting visual information to flow through. This smooth transition from identity to full transformation lets stable optimization, preventing the shock to the system that would occur if visual information were injected at full strength before the model learned how to interpret it.
Complete GATED XATTN-DENSE Block
Now we combine the cross-attention with a feed-forward network to create the full GATED XATTN-DENSE block used in Flamingo.
class GatedXATTNDense(nn.Module):
"""
Complete Flamingo layer: Gated Cross-Attention + Gated Dense (FFN).
Inserted between frozen LM layers.
"""
def __init__(
self,
dim: int = 512,
num_heads: int = 8,
ff_mult: int = 4,
dropout: float = 0.0,
):
super().__init__()
# Gated Cross-Attention component
self.cross_attn = GatedCrossAttention(dim, num_heads, dropout=dropout)
# Gated Dense (FFN) component
self.ffn_norm = nn.LayerNorm(dim)
self.ffn = nn.Sequential(
nn.Linear(dim, dim * ff_mult),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(dim * ff_mult, dim),
nn.Dropout(dropout),
)
# Separate gating parameter for FFN
self.ffn_alpha = nn.Parameter(torch.tensor(0.0))
def forward(
self,
hidden_states: torch.Tensor,
visual_tokens: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Args:
hidden_states: (batch, seq_len, dim) - from previous LM layer
visual_tokens: (batch, num_visual, dim) - from Perceiver Resampler
attention_mask: Optional padding mask
Returns:
updated_hidden: (batch, seq_len, dim)
"""
# Step 1: Gated Cross-Attention to visual tokens
hidden = self.cross_attn(hidden_states, visual_tokens, attention_mask)
# Step 2: Gated Dense (FFN)
normed = self.ffn_norm(hidden)
ffn_out = self.ffn(normed)
ffn_gate = torch.tanh(self.ffn_alpha)
hidden = hidden + ffn_gate * ffn_out
return hiddenIntegration Example
Let's see how these components fit together in a forward pass with interleaved text and images.
class SimpleFlamingoLayer(nn.Module):
"""
Simplified Flamingo-compatible transformer layer.
In practice, this wraps a frozen LM layer with GATED XATTN-DENSE.
"""
def __init__(self, dim: int = 512, num_heads: int = 8):
super().__init__()
# Simulated frozen LM component (in reality, this would be a frozen layer)
self.lm_self_attn = nn.MultiheadAttention(
dim, num_heads, batch_first=True
)
self.lm_ffn = nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, dim * 4),
nn.GELU(),
nn.Linear(dim * 4, dim),
)
# Trainable Flamingo components
self.gated_xattn_dense = GatedXATTNDense(dim, num_heads)
def forward(
self,
x: torch.Tensor,
visual_tokens: Optional[torch.Tensor] = None,
causal_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
# Standard LM self-attention (frozen in real Flamingo)
attn_out, _ = self.lm_self_attn(x, x, x, attn_mask=causal_mask)
x = x + attn_out
# FFN (frozen in real Flamingo)
x = x + self.lm_ffn(x)
# Flamingo cross-attention to images (trainable)
if visual_tokens is not None:
x = self.gated_xattn_dense(x, visual_tokens)
return xText input shape: torch.Size([2, 50, 512]) Visual tokens shape: torch.Size([2, 192, 512]) Output shape: torch.Size([2, 50, 512]) Number of trainable parameters in gated layers: 3,150,850
This shows how visual tokens are integrated into the text processing stream. The text maintains its autoregressive structure (enforced by the causal mask), while cross-attention allows every text position to attend to all visual tokens from preceding images. This architecture lets the model to reference visual information at any point during text generation, whether describing specific details, answering questions about image content, or comparing multiple images in sequence.