tinyvlm-470m

TinyVLM-470M

A 470M parameter Vision-Language Model built from scratch, implementing LLaMA 4-style architecture ideas in a compact multimodal training pipeline.

Model Description

This is an educational implementation of a complete VLM training pipeline, from text pretraining through preference optimization.

Architecture Highlights

Architecture Features:

• Grouped Query Attention (GQA): 12 query heads, 4 KV heads (3x lower KV-cache memory)
• iRoPE: Interleaved RoPE/NoPE layers (3:1 pattern) with chunked attention
• Mixture of Experts: 8 experts, top-2 routing, shared expert
• SwiGLU: Gated activation in all FFN blocks

Training Pipeline

Checkpoints

Usage

# Clone repo and install dependencies
git clone https://github.com/siddhantmedar/tinyvlm-470m
cd tinyvlm-470m
pip install torch torchvision huggingface_hub tokenizers

import torch
import tomllib
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms
from tokenizers import Tokenizer

# Download checkpoints (cached automatically)
llm_ckpt = hf_hub_download("medarsiddhant/tinyvlm-470m", "checkpoints/text_pretraining_best.pt")
vlm_ckpt = hf_hub_download("medarsiddhant/tinyvlm-470m", "checkpoints/vision_alignment_best.pt")
dpo_ckpt = hf_hub_download("medarsiddhant/tinyvlm-470m", "checkpoints/dpo_best.pt")

# Load model
from visual_instruction_tuning.model import create_instruct_vlm

with open("config.toml", "rb") as f:
    config = tomllib.load(f)

model = create_instruct_vlm(config, llm_ckpt, vlm_ckpt)
ckpt = torch.load(dpo_ckpt, map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model_state_dict"])
model = model.to("cuda").eval()

# Load tokenizer
tokenizer = Tokenizer.from_file("visual_instruction_tuning/bpe_tokenizer_with_image_tag.json")
IMAGE_TOKEN_ID = tokenizer.token_to_id("<image>")
EOS_TOKEN_ID = tokenizer.token_to_id("</s>")

# Image transform
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Inference function
@torch.no_grad()
def ask(image_path, question, max_tokens=100):
    image = Image.open(image_path).convert("RGB")
    image_tensor = transform(image).unsqueeze(0).to("cuda")

    prompt = f" USER: {question} ASSISTANT:"
    input_ids = [IMAGE_TOKEN_ID] + tokenizer.encode(prompt).ids
    input_ids = torch.tensor([input_ids], device="cuda")

    generated = []
    for _ in range(max_tokens):
        logits = model(image_tensor, input_ids)
        next_token = logits[0, -1, :].argmax().item()
        if next_token == EOS_TOKEN_ID:
            break
        generated.append(next_token)
        input_ids = torch.cat([input_ids, torch.tensor([[next_token]], device="cuda")], dim=1)

    return tokenizer.decode(generated)

# Example usage
response = ask("path/to/image.jpg", "What is in this image?")
print(response)

Limitations

• Small model size: 470M params limits reasoning capacity compared to 7B+ models
• Repetition: May produce repetitive outputs (use repetition_penalty=1.2)
• Training data: Limited to COCO images and synthetic captions
• Educational purpose: Not intended for production use

Training Hardware

• GPU: NVIDIA RTX 3090 (24GB)
• Total training time: ~35-40 hours across all phases

Links

• Code: GitHub - tinyvlm-470m

Citation

@misc{tinyvlm-470m,
  author = {Siddhant Medar},
  title = {TinyVLM-470M: A Vision-Language Model from Scratch},
  year = {2025},
  publisher = {HuggingFace},
  url = {https://huggingface.co/medarsiddhant/tinyvlm-470m}
}

License

MIT