Vision-Language Projection › Code Implementation

# Configuration mimicking CLIP ViT-L/14 -> LLaMA-2 setup
config = {
    "vision_dim": 1024,  # ViT-L/14 output dimension
    "llm_dim": 4096,  # LLaMA-2-7B embedding dimension
    "num_patches": 256,  # (224/14)^2
    "num_queries": 32,  # Q-Former queries
    "qformer_dim": 768,  # Q-Former hidden dimension
    "qformer_heads": 12,
    "qformer_layers": 12,
}

import torch.nn as nn


class LinearProjection(nn.Module):
    """
    Simple linear projection from vision space to LLM space.
    Maintains sequence length, transforms dimensionality.
    """

    def __init__(self, vision_dim, llm_dim, bias=False):
        super().__init__()
        self.projection = nn.Linear(vision_dim, llm_dim, bias=bias)

    def forward(self, vision_features):
        """
        Args:
            vision_features: [batch, num_patches, vision_dim]
        Returns:
            projected: [batch, num_patches, llm_dim]
        """
        return self.projection(vision_features)

class MLPProjection(nn.Module):
    """
    Two-layer MLP projection with GELU activation.
    Adds non-linear capacity between modalities.
    """

    def __init__(self, vision_dim, llm_dim, hidden_dim=None, dropout=0.1):
        super().__init__()
        if hidden_dim is None:
            hidden_dim = 4 * vision_dim  # Common expansion ratio

        self.fc1 = nn.Linear(vision_dim, hidden_dim)
        self.activation = nn.GELU()
        self.dropout = nn.Dropout(dropout)
        self.fc2 = nn.Linear(hidden_dim, llm_dim)

    def forward(self, vision_features):
        x = self.fc1(vision_features)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

class CrossAttentionLayer(nn.Module):
    """
    Single transformer layer with self-attention on queries
    and cross-attention to vision features.
    """

    def __init__(self, dim, num_heads, vision_dim):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**-0.5

        # Self-attention (queries attend to themselves)
        self.q_self_attn = nn.MultiheadAttention(
            dim, num_heads, batch_first=True
        )
        self.q_self_norm = nn.LayerNorm(dim)

        # Cross-attention (queries attend to vision)
        self.cross_attn = nn.MultiheadAttention(
            dim, num_heads, batch_first=True, kdim=vision_dim, vdim=vision_dim
        )
        self.cross_norm = nn.LayerNorm(dim)

        # FFN
        self.ffn = nn.Sequential(
            nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)
        )
        self.ffn_norm = nn.LayerNorm(dim)

    def forward(self, queries, vision_features):
        # Self-attention with residual
        attn_out, _ = self.q_self_attn(
            queries, queries, queries, need_weights=False
        )
        queries = self.q_self_norm(queries + attn_out)

        # Cross-attention with residual (queries attend to vision)
        cross_out, _ = self.cross_attn(
            queries, vision_features, vision_features, need_weights=False
        )
        queries = self.cross_norm(queries + cross_out)

        # FFN with residual
        ffn_out = self.ffn(queries)
        queries = self.ffn_norm(queries + ffn_out)

        return queries


class QFormer(nn.Module):
    """
    Querying Transformer that compresses variable-length vision features
    into fixed number of learnable query tokens.
    """

    def __init__(self, num_queries, vision_dim, dim, num_heads, num_layers):
        super().__init__()
        self.num_queries = num_queries
        self.dim = dim

        # Learnable query embeddings (the "bottleneck")
        self.query_tokens = nn.Parameter(torch.randn(1, num_queries, dim))

        # Input projection if vision_dim != dim
        self.vision_proj = (
            nn.Linear(vision_dim, dim) if vision_dim != dim else nn.Identity()
        )

        # Transformer layers
        self.layers = nn.ModuleList(
            [
                CrossAttentionLayer(dim, num_heads, dim)
                for _ in range(num_layers)
            ]
        )

        # Final norm
        self.norm = nn.LayerNorm(dim)

    def forward(self, vision_features):
        batch_size = vision_features.size(0)

        # Expand queries for batch
        queries = self.query_tokens.expand(batch_size, -1, -1)

        # Project vision features to match query dimension if needed
        vision_proj = self.vision_proj(vision_features)

        # Pass through transformer layers
        for layer in self.layers:
            queries = layer(queries, vision_proj)

        return self.norm(queries)

class VisionLanguageProjection(nn.Module):
    """
    Complete projection module supporting linear, MLP, and Q-Former approaches.
    """

    def __init__(self, config, method="linear"):
        super().__init__()
        self.method = method
        self.config = config

        if method == "linear":
            self.projector = LinearProjection(
                config["vision_dim"], config["llm_dim"], bias=False
            )
            self.output_tokens = config["num_patches"] + 1  # patches + CLS

        elif method == "mlp":
            self.projector = MLPProjection(
                config["vision_dim"],
                config["llm_dim"],
                hidden_dim=4 * config["vision_dim"],
            )
            self.output_tokens = config["num_patches"] + 1

        elif method == "qformer":
            self.qformer = QFormer(
                num_queries=config["num_queries"],
                vision_dim=config["vision_dim"],
                dim=config["qformer_dim"],
                num_heads=config["qformer_heads"],
                num_layers=config["qformer_layers"],
            )
            # Project from Q-Former dim to LLM dim
            self.llm_proj = nn.Linear(config["qformer_dim"], config["llm_dim"])
            self.output_tokens = config["num_queries"]

        else:
            raise ValueError(f"Unknown projection method: {method}")

    def forward(self, vision_features):
        if self.method in ["linear", "mlp"]:
            return self.projector(vision_features)
        else:  # qformer
            queries = self.qformer(vision_features)
            return self.llm_proj(queries)

def create_vlm_with_projection(vision_encoder, llm, projection_module):
    """
    Assembles a vision-language model with proper freezing.
    Only the projection parameters will be trainable.
    """
    # Freeze vision encoder
    for param in vision_encoder.parameters():
        param.requires_grad = False

    # Freeze LLM (common for initial training)
    for param in llm.parameters():
        param.requires_grad = False

    # Projection remains trainable
    for param in projection_module.parameters():
        param.requires_grad = True

    return nn.ModuleDict(
        {"vision": vision_encoder, "projection": projection_module, "llm": llm}
    )


def training_step_with_masking(
    model, images, text_tokens, text_labels, loss_fn, optimizer
):
    """
    Single training step with proper loss masking.

    text_labels should have -100 for visual token positions
    (ignored by CrossEntropyLoss) and real token ids for text positions.
    """
    model.train()
    optimizer.zero_grad()

    # Forward pass through frozen vision encoder (no grad needed)
    with torch.no_grad():
        vision_features = model["vision"](images)

    # Forward pass through trainable projection
    projected = model["projection"](vision_features)

    # Concatenate visual and text tokens for LLM
    combined_input = torch.cat([projected, text_tokens], dim=1)

    # Build labels: -100 for visual positions (masked), real ids for text
    num_visual = projected.shape[1]
    visual_mask = torch.full(
        (images.shape[0], num_visual),
        fill_value=-100,
        dtype=torch.long,
        device=images.device,
    )
    combined_labels = torch.cat([visual_mask, text_labels], dim=1)

    # Get LLM predictions
    logits = model["llm"](combined_input)

    # Loss computed only on text positions (-100 positions are ignored)
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = combined_labels[..., 1:].contiguous()

    loss = loss_fn(
        shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
    )

    loss.backward()
    optimizer.step()

    return loss.item()