Reconstruction Alignment

Chapter 1: The Reconstruction Objective

The core idea is simple: in addition to the standard next-token prediction loss, add a loss that measures how well the model can reconstruct the input image from its hidden representations.

Standard VLM training (no reconstruction)

The model only needs to predict text. It can ignore any visual information that isn't needed for the text output.

With reconstruction alignment

Where h_image is the model's hidden representation of the image, and Decoder is a lightweight image decoder that maps hidden states back to pixel space (or latent space). The reconstruction loss L_recon penalizes the model if its internal representation has lost visual information.

python
# Standard VLM: text loss only
text_logits = model(image, prompt)
loss = cross_entropy(text_logits, target_text)

# Reconstruction alignment: text + image reconstruction
hidden, text_logits = model(image, prompt, return_hidden=True)
loss_text = cross_entropy(text_logits, target_text)

# Reconstruct image from hidden states
img_hidden = hidden[image_positions]          # [B, N, D]
img_recon = image_decoder(img_hidden)         # [B, C, H, W]
loss_recon = F.mse_loss(img_recon, image_latents)  # reconstruction error

loss = loss_text + lambda_recon * loss_recon

Chapter 2: Architecture

Reconstruction alignment adds a lightweight image decoder to the standard VLM architecture. The decoder maps the model's hidden states at image positions back to pixel space.

The decoder design

The decoder is deliberately lightweight (~10-50M parameters vs the backbone's billions). This is intentional: if the decoder were powerful, it could "cheat" by reconstructing images from minimal information. A weak decoder forces the backbone to do the heavy lifting of preserving visual information.

python
# Lightweight reconstruction decoder
class ReconDecoder(nn.Module):
    def __init__(self, hidden_dim, latent_dim=8):
        self.proj = nn.Linear(hidden_dim, latent_dim * 4 * 4)
        self.upsample = nn.Sequential(
            nn.ConvTranspose2d(latent_dim, 64, 4, 2, 1),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 4, 2, 1),
            nn.ReLU(),
            nn.ConvTranspose2d(32, latent_dim, 4, 2, 1),
        )  # ~10M params total

    def forward(self, hidden_states):
        # hidden_states: [B, N_img, D] from transformer backbone
        h = self.proj(hidden_states)          # [B, N_img, latent*4*4]
        h = h.reshape(-1, 8, 4, 4)           # [B*N, 8, 4, 4]
        recon = self.upsample(h)              # [B*N, 8, 32, 32]
        return recon

Chapter 3: Training Pipeline

Reconstruction alignment can be added to any existing VLM training pipeline. The paper shows it improves both models trained from scratch and existing models fine-tuned with the additional objective.

The λ hyperparameter

The reconstruction weight λ controls how much the model prioritizes visual preservation vs text generation. The paper finds that λ = 0.1 works well across model sizes — enough to improve visual representations without degrading text quality.

Chapter 4: Why It Works

Reconstruction alignment works through an information bottleneck mechanism. The key insight is about what information the model chooses to preserve in its hidden states.

Without reconstruction: lossy encoding

Standard VLMs encode images through a vision encoder (CLIP ViT) and project them into the LM's space. The LM then generates text. The LM only needs to preserve visual information that helps predict the next text token. Fine-grained details (exact colors, spatial arrangements, small objects) are often discarded because they're not needed for generating a correct caption.

With reconstruction: lossless encoding

When the model must also reconstruct the image, it can't afford to discard any visual details. The hidden states must contain enough information to regenerate the full image. This forces richer, more detailed visual representations that contain fine-grained information the text objective alone wouldn't preserve.

Chapter 5: Ablations

The paper provides detailed ablations showing which design choices matter most.

Chapter 6: Results & Showcase

Reconstruction alignment consistently improves multimodal understanding benchmarks across model sizes.

Chapter 7: Connections

Reconstruction alignment connects to a deep principle in representation learning: the best representations are those that preserve the most information about the input while remaining useful for downstream tasks.