LLaVA Architecture › Two-Stage Training Strategy

Michael BrenndoerferFebruary 11, 20269 min read
Language AI HandbookAll chapters
Part 50 of 67Multimodal Models
Chapter 385LLaVA Architecture › Two-Stage Training Strategy
Reading section 4 of 14 · Two-Stage Training Strategy · 9 minSection 4 of 14 · 9 min

Book contents

Section 4 of 14 · 9 min read

LLaVA Architecture › Two-Stage Training StrategyOne section · Sections in this chapter4 / 14 ↓
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

Training a multimodal model end-to-end from random initialization would require enormous datasets and computational resources, likely billions of image-text pairs and thousands of GPU hours. Such training would also risk destabilizing the carefully learned representations of pre-trained components. LLaVA's breakthrough insight is decomposing the training into two distinct stages, each with specific objectives and data requirements. This mirrors the progressive training approach discussed in earlier instruction tuning chapters, but extends it to the multimodal domain. The curriculum follows a pedagogical intuition: first teach the model to recognize what it sees, then teach it to converse about what it sees.

Stage 1: Feature Alignment Pre-training

The first stage focuses exclusively on teaching the projection layer to translate visual features into language model embeddings. This stage is analogous to teaching a translator the basic vocabulary of a new language before asking them to interpret complex poetry. During this phase:

  • The vision encoder remains frozen
  • The language model remains frozen
  • Only the projection matrix (or MLP) is trained

Training Data: This stage uses image-caption pairs from datasets like CC3M (Conceptual Captions), filtered to approximately 600,000 high-quality pairs. Unlike complex instruction data, these are simple descriptions: "a dog playing in the park" or "a red apple on a wooden table." These captions give direct, unambiguous mappings between visual content and linguistic descriptions. The simplicity is important: we want the projection layer to learn that a specific constellation of visual features corresponds to the concept "dog" before attempting to handle ambiguous or complex instructions.

The filtering step that produces 600K from the full 3.3M CC3M pairs is important. Raw web-scraped image captions often contain noise, off-topic text (like news article text where the image is incidental), or very short uninformative descriptions. Filtering by caption length, image quality, and relevance ensures the projection layer trains on clean signal. The selected pairs tend to have captions that accurately describe what is visually present, which is precisely the relationship the projection layer needs to learn.

Objective: The model learns to predict the caption text given the image. We minimize the cross-entropy loss between the predicted probability distribution over the vocabulary and the actual caption tokens. For a sequence of LL tokens, the loss is calculated by summing the negative log-likelihoods of each correct token:

L=−∑i=1Llog⁡P(ti∣t<i,Hv)\mathcal{L} = -\sum_{i=1}^{L} \log P(t_i \mid t_{<i}, H_v)

where:

  • L\mathcal{L}: the total cross-entropy loss across the caption
  • LL: the number of tokens in the target caption
  • tit_i: the ii-th target token in the caption (the token we want to predict)
  • t<it_{<i}: all tokens preceding position ii in the sequence (the context for prediction)
  • P(ti∣t<i,Hv)P(t_i \mid t_{<i}, H_v): the probability assigned by the model to the correct token tit_i, conditioned on previous tokens t<it_{<i} and projected visual features HvH_v
  • HvH_v: the projected visual embeddings, serving as the visual context for the language model

This stage effectively aligns the visual and linguistic representations. The projection layer learns that certain visual patterns should map to specific semantic concepts in the language model's embedding space. Because the LLM is frozen, the projection layer must adapt to the LLM's existing semantic structure rather than vice versa. This constraint forces the projection to discover the "language" that the LLM already understands, effectively creating a visual vocabulary comprehensible to the text model.

During this stage, the projection layer is performing a form of reverse-engineering the LLM's embedding space. Given that the LLM knows what "dog" means, the projection layer must find the direction in RdLLM\mathbb{R}^{d_{\text{LLM}}} that activates the LLM's dog-related representations, and it must map CLIP's visual representation of a dog there. This is a well-defined optimization problem with a clear signal, and the 600K training pairs give enough variation (different dog breeds, lighting conditions, contexts) to learn a general mapping rather than an overfitted one.

Why Freeze the LLM?

Freezing the language model during alignment serves two purposes. First, it preserves the LLM's general knowledge and reasoning capabilities, preventing catastrophic forgetting during multimodal adaptation. Second, it constrains the projection layer to learn representations compatible with the LLM's existing embedding space. This keeps visual concepts are expressed in terms the LLM already understands. If both the LLM and the projection layer changed simultaneously, they could reach a local optimum where the LLM has adapted to whatever the projection layer outputs rather than learning to process semantically real visual features.

Stage 2: Visual Instruction Tuning

Once the projection layer has learned basic vision-language alignment, the second stage teaches the model to follow complex, multimodal instructions. This stage adapts the instruction tuning paradigm to the visual domain. While Stage 1 teaches the model to describe, Stage 2 teaches it to converse and reason before following directions.

Modifications from Stage 1:

  • The vision encoder remains frozen
  • The projection layer continues training
  • The language model is unfrozen and trained

Unfreezing the language model in Stage 2 is needed because caption prediction and instruction following require quite different behaviors. During Stage 1, the model only needs to output descriptions of image content. During Stage 2, it needs to answer questions, explain reasoning, refuse unreasonable requests, and adapt its response style to different instruction types. These higher-order behaviors require the language model's weights to update, incorporating visual context into the patterns that govern response generation.

Training Data: This stage requires instruction-following data with visual inputs. LLaVA pioneered the creation of such data using GPT-4V (or text-only GPT-4 with image captions) to generate diverse instruction-response pairs. The dataset includes:

  • Conversational data: Multi-turn dialogues about images ("What color is the car? Is it moving?")
  • Detailed description: Elaborate image captions with spatial relationships
  • Complex reasoning: Visual question answering requiring inference ("Why might this person be carrying an umbrella?")

The data format follows standard chat templates:

Human: <image>\nDescribe what you see in this image. Assistant: The image shows a bustling farmers market with...

Loss Masking: Following best practices from instruction tuning, the loss is computed only on the assistant's responses, not on the user's questions or the image tokens. This ensures the model learns to generate appropriate answers rather than simply memorizing the conversation structure. Without this masking, the model might learn to parrot back the question format or hallucinate image features, whereas masked training forces it to focus on creating relevant, grounded responses.

The two-stage approach creates a curriculum: first learn to see, then learn to converse about what you see. This separation prevents the instability that would arise from trying to learn both alignment and instruction following simultaneously. Attempting both at once would create conflicting gradients: the projection layer would receive mixed signals about whether to prioritize simple feature matching or complex reasoning, likely resulting in poor performance on both objectives.

Out[4]:
Visualization
Stacked bar chart comparing trainable parameter counts across Stage 1 and Stage 2 of LLaVA training.
Trainable parameter counts in each stage of LLaVA training. Stage 1 trains only the lightweight projection layer (~16M parameters), while Stage 2 additionally opens up the language model (~7B parameters) for fine-tuning. The vision encoder remains frozen throughout both stages, preserving its CLIP representations.

Compute Efficiency of the Two-Stage Design

The two-stage design has a remarkable compute advantage. Stage 1, which requires training only the projection layer, converges in roughly 4 hours on 8 A100 GPUs with the 600K CC3M subset. The projection layer has only about 16 million parameters for the ViT-L to LLaMA-7B setting, making each gradient update extremely cheap. By contrast, Stage 2 trains the full 7 billion parameter LLM, but uses only 150K instruction-following samples rather than billions of image-text pairs, keeping the total compute comparable to the fine-tuning budget for a text-only instruction tuned model.

This frugality is one of LLaVA's most important contributions. The original paper demonstrated that a single 8-GPU machine could complete the entire training pipeline in under a day, putting multimodal model training within reach of university research groups and small companies. Prior work required clusters of hundreds of GPUs and weeks of training. The democratizing effect of this efficiency has been significant: dozens of domain-specific and multilingual variants of LLaVA have been developed by groups that would not have had the resources to train a vision-language model from scratch.