LMFusion: Adapting LLMs for Multimodal

Chapter 1: Why Not From Scratch?

Training a multimodal model from scratch (Chameleon, Transfusion) requires relearning everything the LLM already knows. The compute cost is enormous:

That's roughly a 9x cost reduction. And the pretrained LLM's text abilities are typically better than what you'd get from mixed-modal training (because it was trained on more text data and didn't have to "share" capacity with images).

The catastrophic forgetting problem

Simply fine-tuning a pretrained LLM on image data destroys its text abilities. The model's weights shift to accommodate image patterns, overwriting the carefully learned text representations. This is catastrophic forgetting — a well-known problem in continual learning.

python
# Naive adaptation: catastrophic forgetting
model = load_pretrained("llama-3-8b")
for batch in image_data:
    loss = model(batch)   # training on images
    loss.backward()       # ALL weights updated
    optimizer.step()      # text knowledge destroyed!

# LMFusion: careful adaptation
model = load_pretrained("llama-3-8b")
model.freeze(model.text_ffn)      # protect text-specific params
model.add(image_input_proj)        # new image params
model.add(image_output_head)       # new image params
# Only update: attention (shared) + new image params

Chapter 2: Architecture Design

LMFusion's architecture extends the Transfusion paradigm (AR text + diffusion images) but applied to an existing LLM backbone. The key design decisions determine which parameters are new, which are shared, and which are frozen.

Parameter categories

python
# LMFusion layer structure
class LMFusionLayer(nn.Module):
    def __init__(self, pretrained_layer, config):
        # Reuse pretrained attention (will be fine-tuned)
        self.attn = pretrained_layer.attention  # trainable

        # Keep text FFN frozen
        self.text_ffn = pretrained_layer.ffn    # FROZEN
        self.text_ffn.requires_grad_(False)

        # Add NEW image FFN (randomly initialized)
        self.image_ffn = SwiGLU_FFN(config.dim, config.dim * 4)  # trainable
        self.image_norm = RMSNorm(config.dim)  # trainable

    def forward(self, x, modality):
        # Shared attention (fine-tuned)
        h = self.attn(x) + x

        # Route to modality-specific FFN
        out = torch.empty_like(h)
        out[modality == 0] = self.text_ffn(h[modality == 0]) + h[modality == 0]    # frozen
        out[modality == 1] = self.image_ffn(self.image_norm(h[modality == 1])) + h[modality == 1]  # trained
        return out

Chapter 3: Training Strategy

LMFusion uses a multi-stage training strategy that gradually introduces image capability while monitoring text preservation.

Chapter 4: Progressive Unfreezing

LMFusion doesn't unfreeze all attention layers at once. It uses progressive unfreezing: first the top layers (closest to output), then gradually lower layers. This protects the low-level representations that are most important for text quality.

Why does layer order matter? In transformers, lower layers learn general features (syntax, token relationships) while upper layers learn task-specific features. Image generation mostly needs task-specific adaptations (upper layers), while the general features (lower layers) should be preserved.

Chapter 5: Text Preservation

LMFusion monitors text quality throughout training and introduces several techniques to prevent degradation beyond acceptable thresholds.

Data mixing during adaptation

Even during image training, LMFusion includes text-only data in every batch (typically 30-50% text). This acts as a regularizer, continuously reminding the model of text patterns and preventing drift.

Text quality metrics throughout training

Chapter 6: Results & Showcase

LMFusion demonstrates that adapted LLMs can compete with models trained from scratch, at a fraction of the cost.

Chapter 7: Connections

LMFusion demonstrates that the future of multimodal AI may not require training everything from scratch. Instead, we can build on the massive investment in text-only LLMs.