SwinginBaloo/sdpo-fewshot-ocr-qwen2vl-2b

SDPO Few-Shot Bank Integrator for VLM Document Extraction

This project implements Self-Distillation Policy Optimization (SDPO) to train a small VLM (Qwen2-VL-2B-Instruct) to perform document information extraction without few-shot examples, by distilling knowledge from its own few-shot-augmented teacher.

Paper Reference

• SDPO: arXiv:2601.20802 — Self-Distillation Policy Optimization

Approach

Teacher vs Student

• Teacher: The same model prompted with few-shot examples from a dynamic bank + the document image + the extraction question. This provides rich feedback/context.
• Student: The same model prompted with ONLY the document image + the extraction question (no few-shot examples).
• Goal: Distill the few-shot capability into zero-shot weights via SDPO.

Algorithm (SDPO for Document Extraction)

1. Generate completions from the student (on-policy)
2. Compute exact-match reward against ground truth answers
3. Compute log-probs of the SAME completion under:
   • Teacher prompt (with few-shot examples) — stop-grad
   • Student prompt (without examples) — with grad
4. Compute SDPO advantage: A_t = log P_teacher(y_t) - log P_student(y_t)
   • Positive advantage where teacher is more confident → reinforce student
   • Negative advantage where teacher disagrees → penalize student
5. Scale advantages by reward signal and compute policy gradient loss
6. Repeat — the student gradually internalizes the few-shot reasoning

Dataset

• DocVQA (nielsr/docvqa_1200_examples) — document images with questions and answers
• Few-shot bank: first 25 examples per split
• Training: remaining examples with random few-shot sampling from bank

Model

• Qwen/Qwen2-VL-2B-Instruct (~2B parameters, vision-language model)
• Fits on a T4 GPU with bfloat16 + gradient checkpointing

Files

• prepare_sdpo_dataset.py — prepares DocVQA with student/teacher prompts
• train_sdpo_vlm.py — SDPO training loop

Usage

# Prepare dataset (creates docvqa_sdpo_train/ and docvqa_sdpo_test/)
python prepare_sdpo_dataset.py

# Train (requires CUDA GPU, ~16GB VRAM)
python train_sdpo_vlm.py

Hyperparameters

Param	Value
Learning rate	5e-6
Epochs	2
Micro batch size	1
Gradient accumulation	8 (effective batch = 8)
Max completion length	64 tokens
Temperature	0.7
Top-p	0.9
Max examples	600
Few-shot bank size	25
Few-shot per sample	2

Why This Works for Continual Learning

The few-shot bank acts as an episodic memory. During training, we randomly sample different subsets of examples for each training step. The model learns to extract patterns from these examples and internalize them into its weights. Over time, the zero-shot student matches the few-shot teacher, effectively "learning" from the bank without needing the bank at inference time. New examples can be continually added to the bank and the model re-trained.

Generated by ML Intern

This model repository was generated by ML Intern, an agent for machine learning research and development on the Hugging Face Hub.

• Try ML Intern: https://smolagents-ml-intern.hf.space
• Source code: https://github.com/huggingface/ml-intern