You can deploy this model and start earning money today!

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

Name: sdpo-fewshot-ocr-qwen2vl-2b
Availability: InStock

ml-intern

2601.20802

Views

No views yet

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)

Generate completions from the student (on-policy)
Compute exact-match reward against ground truth answers
Compute log-probs of the SAME completion under:
- Teacher prompt (with few-shot examples) — stop-grad
- Student prompt (without examples) — with grad
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
Scale advantages by reward signal and compute policy gradient loss
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

bash

# 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

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