Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

LLaVA Architecture › Training Dynamics and Loss Computation

Michael BrenndoerferFebruary 11, 20267 min read

Language AI HandbookAll chapters

Part 50 of 67Multimodal Models

Chapter 385LLaVA Architecture › Training Dynamics and Loss Computation

Reading section 10 of 14 · Training Dynamics and Loss Computation · 7 minSection 10 of 14 · 7 min

Section 10 of 14 · 7 min read

LLaVA Architecture › Training Dynamics and Loss ComputationOne section · Sections in this chapter10 / 14 ↓

Sections in this chapter14

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

Understanding how LLaVA trains requires examining the specific loss computation, particularly the handling of multimodal sequences and label masking. The training process must carefully balance supervision across different parts of the input to prevent the model from learning spurious correlations or trivial solutions.

Sequence ConstructionLink Copied

A training example in LLaVA consists of a sequence combining visual markers, visual features, and text tokens:

S = [\text{<image>}, V_1, V_2, \ldots, V_N, \text{<text\_tokens>}]

where:

$S$ : the complete multimodal input sequence
$\text{<image>}$ : a special token marking the image position (often replaced by projected visual features)
$V_1, \ldots, V_N$ : the projected visual tokens representing image patches (typically 576 tokens for 336px resolution)
$\text{<text\_tokens>}$ : the sequence of tokenized instruction and response text

The full input to the LLM combines the projected visual features with embedded text tokens:

\text{Inputs} = [\text{Projected}(Z);\, \text{Embed}(T_{\text{instruction}});\, \text{Embed}(T_{\text{response}})]

where:

$\text{Projected}(Z)$ : the visual features after transformation through the projection layer
$T_{\text{instruction}}$ : the token IDs representing the user's question or prompt text
$T_{\text{response}}$ : the token IDs representing the assistant's target answer
$\text{Embed}(\cdot)$ : the language model's embedding lookup function converting token IDs to dense vectors
$\text{Inputs}$ : the final concatenated sequence of embeddings fed to the transformer layers

This concatenation treats visual tokens and text tokens identically once they enter the transformer, letting the self-attention mechanism to establish relationships between any visual patch and any text token without architectural constraints on cross-modal interaction. There is no separate "image processing" module inside the transformer: visual tokens and text tokens flow through the same attention heads, the same feed-forward networks, and the same layer normalization operations.

Loss Masking StrategyLink Copied

As discussed in earlier chapters on instruction tuning, we compute loss only on the tokens we want the model to learn to generate. For LLaVA:

Image tokens: No loss computed (the model does not predict these)
User instruction tokens: No loss computed (we do not train the model to predict the question)
Assistant response tokens: Cross-entropy loss computed
Padding tokens: Ignored in loss computation

This masking ensures the model learns to generate appropriate responses conditioned on images and questions, rather than memorizing the questions themselves or trying to reconstruct image features. Training on user questions would encourage the model to parrot back query patterns rather than learn to answer them, while training on image tokens would waste capacity on reconstructing visual inputs that are already provided directly.

The loss at each unmasked position is:

\mathcal{L}_i = -\log P(t_i \mid t_{<i}, H_v; \theta)

where $\theta$ is all trainable parameters. The total loss averages over unmasked positions:

\mathcal{L} = \frac{1}{|\mathcal{A}|} \sum_{i \in \mathcal{A}} \mathcal{L}_i

where $\mathcal{A}$ denotes the set of active (non-masked) token positions. Averaging over active tokens rather than all positions ensures that examples with short responses receive the same effective learning rate per response token as examples with long responses, preventing the model from implicitly down-weighting short answers.

In[13]:

Code

import torch.nn as nn


def compute_llava_loss(logits, labels, vision_token_len=576):
    """
    Compute masked loss for LLaVA training.
    Labels should be -100 for tokens to ignore.
    """
    # Shift logits and labels for next-token prediction
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()

    # Flatten for cross entropy
    loss_fct = nn.CrossEntropyLoss()
    shift_logits = shift_logits.view(-1, shift_logits.size(-1))
    shift_labels = shift_labels.view(-1)

    # Compute loss only on non-masked tokens (CrossEntropyLoss ignores -100 by default)
    loss = loss_fct(shift_logits, shift_labels)
    return loss


def create_multimodal_labels(
    input_ids, response_start_idx, vision_token_len=576, pad_token_id=0
):
    """
    Create labels where only response tokens are supervised.
    Prefix (image + question) gets -100 (ignored).
    """
    labels = input_ids.clone()

    # Mask image tokens
    labels[:, :vision_token_len] = -100

    # Mask user instruction tokens up to response start
    labels[:, vision_token_len:response_start_idx] = -100

    # Mask padding tokens
    labels[labels == pad_token_id] = -100

    return labels

The masking strategy is important for stable training. Without proper masking, the model might learn to generate repetitive question-answer pairs instead of learning to respond to novel queries. The use of -100 as the ignore index is standard in PyTorch's CrossEntropyLoss, which skips gradient computation for positions labeled -100 without requiring a separate mask multiplication step. This efficiency detail means the masking adds negligible overhead to training despite eliminating the majority of positions from the gradient computation.

Out[14]:

Visualization

Token sequence diagram showing masked and active loss regions for LLaVA training across image, user, assistant, and padding segments. — Loss masking pattern in LLaVA visual instruction tuning. Image tokens and user instruction tokens contribute no gradient (gray regions), while only assistant response tokens are supervised with cross-entropy loss (blue region). This design teaches the model to generate grounded responses rather than memorizing conversation structure.

Gradient Flow and OptimizationLink Copied

During Stage 2, gradients flow through the entire system except the frozen vision encoder. Starting from the cross-entropy loss on assistant response tokens, gradients propagate backward through the LLM's transformer layers, through the projection layer, and terminate at the projection layer's inputs (which are the frozen vision encoder outputs). This means the vision encoder's parameters never receive updates, but the projection layer learns from the signal that originates in the LLM's loss computation.

This gradient flow has an important implication for what the projection layer learns to stand for. In Stage 1, the projection learns to map visual features to language space in a way that produces correct captions, a relatively local objective. In Stage 2, the projection layer receives gradients that reflect whether its output contributed to correct instruction-following behavior, a much richer and more demanding objective. The Stage 2 gradient signal teaches the projection to produce semantically accurate visual tokens that are useful for the LLM's reasoning processes.

The learning rates used in each stage reflect this difference in objectives. Stage 1 typically uses a higher learning rate (around 1e-3) since the projection layer starts from random initialization and must learn rapidly from a relatively simple objective. Stage 2 uses a much lower learning rate for both the projection layer (around 2e-5) and the LLM (1e-4 to 2e-5 depending on scale), as both are making incremental refinements rather than large initial adjustments.