LLaVA Architecture › Implementation and Usage
Language AI HandbookAll chapters LLaVA Architecture › Implementation and UsageOne section · Sections in this chapter8 / 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
Let us examine how LLaVA functions in practice, from processing an image through generating a response. We'll implement a simplified version of the forward pass to understand the data flow.
First, we establish the setup constants and define the projection layer. In practice, these would be loaded from pre-trained LLaVA weights, but defining them explicitly clarifies the architecture.
import torch.nn as nn
# Configuration matching LLaVA-1.5-7B style architecture
vision_hidden_size = 1024 # CLIP ViT-L
llm_hidden_size = 4096 # LLaMA-7B hidden dimension
num_patches = 576 # 24x24 for 336px with 14px patches
vocab_size = 32000 # LLaMA vocabulary
class LLaVAProjection(nn.Module):
"""MLP connector used in LLaVA 1.5"""
def __init__(self, vision_dim, llm_dim):
super().__init__()
self.fc1 = nn.Linear(vision_dim, llm_dim)
self.act = nn.GELU()
self.fc2 = nn.Linear(llm_dim, llm_dim)
def forward(self, vision_features):
x = self.fc1(vision_features)
x = self.act(x)
x = self.fc2(x)
return xThe projection layer is remarkably simple: just two linear transformations with a GELU activation. This small module (approximately 16M parameters for 1024 to 4096 to 4096) is the only trainable component during the first stage of training. Its compactness is what makes Stage 1 so fast to train, despite working with visual features from a 300M-parameter vision encoder.
class SimpleLLaVA(nn.Module):
"""Simplified LLaVA architecture showing component interaction"""
def __init__(self, vision_model_name="openai/clip-vit-large-patch14-336"):
super().__init__()
# Frozen vision encoder
self.vision_encoder = CLIPVisionModel.from_pretrained(vision_model_name)
for param in self.vision_encoder.parameters():
param.requires_grad = False
# Projection layer (trainable)
self.projector = LLaVAProjection(vision_hidden_size, llm_hidden_size)
self.llm_hidden_size = llm_hidden_size
def encode_image(self, pixel_values):
"""Process image through vision encoder and projection"""
vision_outputs = self.vision_encoder(pixel_values)
# Shape: [batch, num_patches, vision_hidden_size]
# Skip CLS token (index 0), use patch embeddings only
vision_features = vision_outputs.last_hidden_state[:, 1:, :]
# Project to LLM space
projected = self.projector(vision_features)
return projectedThe encode_image method processes images through the frozen CLIP encoder and projects the features. The [:, 1:, :] indexing skips the CLS token at position 0. CLIP's CLS token is a global image summary, but LLaVA uses the individual patch embeddings to preserve spatial information. Each patch embedding captures a specific region of the image, which allows the LLM's attention to relate individual image regions to specific parts of the text query.
import torch
# Demonstrate tensor shapes through the pipeline
batch_size = 2
vision_shape = torch.Size([batch_size, num_patches, vision_hidden_size])
projected_shape = torch.Size([batch_size, num_patches, llm_hidden_size])
num_tokens = projected_shape[1]
dummy_pixel_values = torch.randn(batch_size, 3, 336, 336)Input image shape: torch.Size([2, 3, 336, 336]) After vision encoder: torch.Size([2, 576, 1024]) After projection: torch.Size([2, 576, 4096]) Visual tokens per image: 576
The tensor flow shows the transformation: images start as 3-channel pixel tensors (batch, 3, 336, 336), become 576 visual tokens in 1024-dimensional CLIP space, then project to 4096-dimensional space matching the LLM's embeddings. Each of the 576 patch tokens corresponds to a 14px by 14px region of the original image.
