Visual Instruction Tuning: Where the Loss Function Actually Applies

This article takes apart the engineering mechanics of Multimodal Large Language Models (MLLMs). Today, we’re unpacking the exact mechanics of Visual Instruction Tuning, focusing heavily on how we construct the loss function when mixing image tokens with text tokens.

The Core Concept: Multimodal Auto-Regressive Next-Token Prediction

In standard LLMs, we predict the next token based on all previous tokens. In MLLMs like LLaVA or Qwen-VL, we inject continuous visual features (projected from a Vision Encoder like CLIP) directly into the LLM’s embedding space. To the LLM, an image is just a sequence of special tokens.

The Math: Multi-Modal Auto-Regressive Loss

Given an image $\mathbf{X}_v$ and a text instruction-response sequence $\mathbf{X}_t$, we convert the image into a sequence of tokens $\mathbf{H}_v$ and the text into tokens $\mathbf{H}_t$. The combined sequence is $\mathbf{H} = [\mathbf{H}_v, \mathbf{H}_t]$.

The standard auto-regressive language modeling objective maximizes the probability of generating the target sequence $\mathbf{X}_{response}$ given the instruction $\mathbf{X}_{instruct}$ and image $\mathbf{X}_v$. The loss function $L$ is the negative log-likelihood:

$$L(\theta) = – \sum_{i=1}^{L} \log p_\theta(x_i | \mathbf{X}_v, x_{

Crucial Detail: The Mask Matrix
We do not compute loss on the image tokens, nor do we compute loss on the user’s instruction. We only compute loss on the model’s generated response. Thus, we introduce a boolean mask $\mathbf{M}$, where $m_i = 1$ if the $i$-th token belongs to the target response, and $m_i = 0$ if it belongs to the image or the system/user instruction.

The modified loss becomes:

$$L(\theta) = – \frac{1}{\sum_{i=1}^{L} m_i} \sum_{i=1}^{L} m_i \log p_\theta(x_i | \mathbf{X}_v, x_{

Data Flow: From JSON to Masked Labels

How does the data actually flow from disk to the loss function? Here is the mental model:

graph TD
    A[Multimodal JSON Dataset] -->|Extract| B(Image Path & Conversation)
    B --> C{Tokenizer & Processor}
    C -->|Image| D[Vision Encoder -> Image Tokens]
    C -->|Text| E[LLM Tokenizer -> Text Tokens]
    D --> F[Sequence Concatenation]
    E --> F
    F --> G[Generate Label Mask]
    G -->|Image & Instruction Tokens| H[Label = -100 Ignored]
    G -->|Response Tokens| I[Label = Token ID]
    H --> J[PyTorch CrossEntropyLoss]
    I --> J

PyTorch Implementation: Masking the Labels

When implementing the DataLoader, setting the correct labels is the most bug-prone step. In PyTorch’s CrossEntropyLoss, a label of -100 is ignored during loss computation. We must mask out the image tokens and the prompt tokens.

import torch
from torch.utils.data import Dataset

IGNORE_INDEX = -100
IMAGE_TOKEN_ID = 32000 # Example ID for <image>

class MultimodalDataset(Dataset):
    def __init__(self, data, tokenizer):
        self.data = data
        self.tokenizer = tokenizer
        
    def __getitem__(self, idx):
        item = self.data[idx]
        # Example format: "<image>\nUser: What is this?\nAssistant: A cat."
        text = item['conversations'] 
        
        # 1. Tokenize the entire sequence
        input_ids = self.tokenizer(text, return_tensors="pt").input_ids[0]
        
        # 2. Clone input_ids to create labels
        labels = input_ids.clone()
        
        # 3. Mask out the image tokens
        labels[labels == IMAGE_TOKEN_ID] = IGNORE_INDEX
        
        # 4. Mask out the user instruction (simplified logic for demonstration)
        # Find the delimiter where the assistant's response begins
        assistant_token_id = self.tokenizer.convert_tokens_to_ids("Assistant:")
        try:
            assistant_idx = (input_ids == assistant_token_id).nonzero(as_tuple=True)[0][0]
            # Mask everything up to the assistant's response (including the prompt)
            labels[:assistant_idx + 1] = IGNORE_INDEX
        except IndexError:
            pass # Handle edge cases
            
        return {
            "input_ids": input_ids,
            "labels": labels,
            "image_path": item['image']
        }

Why training runs in two stages, and what stays frozen

Masking decides which positions produce gradients. An equally important question follows: which parameters those gradients are allowed to reach. Visual instruction tuning normally runs in two stages with entirely different freezing policies, and reversing them makes training unstable.

Stage one trains only the projector, with the vision encoder and the LLM both frozen. At this point the projector is randomly initialised, so the visual tokens it emits look like noise to the LLM. Unfreezing the LLM here means gradients pull it toward accommodating that noise — noise which changes at the next step, so the LLM is chasing a moving target. That shows up as slow loss decrease accompanied by degradation of the model’s existing language ability.

Freezing the LLM first gives the projector a fixed target space to align into. This stage needs only large volumes of image-caption pairs; instruction data is not required.

Stage two unfreezes the LLM and fine-tunes on genuine instruction data. By now the projector emits reasonable visual tokens, so the LLM faces a broadly stable input distribution and fine-tuning is meaningful.

The vision encoder normally stays frozen throughout. Its representations were pretrained on enormous image corpora, an instruction dataset is nowhere near large enough to improve them, and unfreezing simply overfits it to the instruction set’s image distribution — presenting as strong performance on training-set-style images and collapse on anything else.

# Stage one: only the projector is trainable
for p in vision_encoder.parameters(): p.requires_grad = False
for p in llm.parameters():            p.requires_grad = False
for p in projector.parameters():      p.requires_grad = True

# Stage two: unfreeze the LLM, vision encoder still frozen
for p in llm.parameters():            p.requires_grad = True

One check that is easy to miss: after setting requires_grad, confirm the optimiser only received trainable parameters. Constructing it with model.parameters() means frozen parameters carry no gradient but an optimiser with weight decay still updates them — the freeze is void, and nothing reports an error.

opt = torch.optim.AdamW(
    [p for p in model.parameters() if p.requires_grad],  # filter is mandatory
    lr=2e-5, weight_decay=0.0)

When the mask is wrong, the loss curve still looks fine

The awkward property of this code is that a wrong mask still trains and the loss still falls. No exception, no error — the model simply learns the wrong thing.

The two most common mistakes each have a signature worth remembering separately.

Forgetting to mask the visual token positions. If labels at visual token positions are not set to the ignore value, the model is asked to “predict the next visual token.” Visual tokens are continuous projector outputs and belong to no vocabulary, so the task is meaningless. The signature is an abnormally high starting loss — a large batch of impossible predictions — that descends to some value and then plateaus. That plateau is the fixed loss contributed by the meaningless portion.

Forgetting to mask the instruction. If the user’s question also contributes to the loss, the model learns to generate plausible user questions alongside answering them. The signature is subtler: the loss curve looks perfectly normal, but at inference the model tends to ask and answer itself — finishing a reply and then inventing the next question.

So the loss curve is not sufficient for acceptance. The reliable approach is to inspect one sample’s label tensor directly:

ignored = (labels == -100).sum().item()
total   = labels.numel()
print(f"ignored {ignored}/{total}")
# decode the part that actually contributes to loss and read it
kept = input_ids[labels != -100]
print(tokenizer.decode(kept))

Print that last line and read it: it should contain only the assistant’s reply — no system prompt, no user question, no image placeholders. The check takes under a minute and eliminates this entire class of problem, whereas judging by loss curve usually means discovering it after training completes and inference behaves strangely.

Production Pitfalls

1. Catastrophic Forgetting

When tuning an MLLM, if the learning rate is too high or the dataset heavily biases towards short visual descriptions, the underlying LLM can quickly lose its general reasoning or conversational capabilities—a phenomenon known as catastrophic forgetting. To mitigate this, practitioners mix text-only instruction tuning data (like Alpaca or ShareGPT datasets) with the visual data to anchor the language model’s capabilities.

2. Why We Freeze the Vision Encoder

In most Visual Instruction Tuning phases (like LLaVA stage 2), we completely freeze the Vision Encoder (e.g., CLIP ViT) and only train the projection layer and the LLM backbone. Why?

• Representation Stability: CLIP is already trained on billions of image-text pairs. It possesses an excellent, aligned manifold of visual concepts. Unfreezing it on a small instruction dataset (e.g., 150k examples) will overfit the encoder and destroy the generalized visual representations.
• Compute Efficiency: ViTs are memory-intensive. Freezing the vision encoder saves VRAM for training the much larger LLM backbone (e.g., 7B or 13B parameters) and allows for higher batch sizes or gradient accumulation steps.

Takeaway: Visual instruction tuning isn’t about teaching the model to “see”—the Vision Encoder already sees. It’s about teaching the LLM to attend to and reason over the continuous visual features that have been injected into its context window, while strictly calculating loss only on the generated textual response.