Introduction

We propose a novel uncertainty-aware knowledge distillation method, which can improve the performance of the student model by leveraging the uncertainty of the teacher model. [Paper]

📸 Release

• Apri. 27th, 2026: Our Beta-KD weights are uploaded on the HuggingFace website. We also provide inference examples so that anyone can enjoy them early.
• Mar. 22th, 2026: The training and evaluation codes of Beta-KD are available now! Follow these step-by-step instructions below to easily train your own Beta-KD in 5 hours ⚡️ !
• Mar. 21th, 2026: 🔥🔥🔥 We release Beta-KD: Uncertainty-Aware Knowledge Distillation for Multimodal Large Language Models on arxiv. Refer to our paper for more details !

🦙 Model Zoo

Model Zoo

🛠️ Install

Clone this repository and install conda environment

git clone git@github.com:Jingchensun/beta-kd.git
cd beta-kd
 
conda create -n beta-kd python=3.10 -y
conda activate beta-kd
pip install --upgrade pip
pip install -r requirements.txt

Step-by-step Tutorial

1. Prepare Data

• For convenience, assume your working directory /path/to/project/Beta-KD as work_dir:

  • cd ${work_dir} && mkdir -p data/pretrain_data data/finetune_data data/benchmark_data

• prepare pre-training data

  • cd ${work_dir}/data/pretrain_data
  • download the ShareGPT4V-PT from here, which is provided by ShareGPT4V team.

• prepare multi-task training data

  • cd ${work_dir}/data/finetune_data
  • download the annotation of our Beta-KD_V2_FT_Mix2M data from huggingface here, and download the images from constituting datasets: Text-VQA, IConQA, SQA, SBU, follow ShareGPT4V to download images from: LAION-CC-SBU-558K, COCO, WebData, SAM, GQA, OCR-VQA, TextVQA, VisualGnome (Part1, Part2)

• prepare evaluation benchmark data

  • We evaluate models on a diverse set of 6 benchmarks, i.e. GQA, MMBench, MME, POPE, SQA, TextVQA. We do not evaluate using beam search to make the inference process consistent with the chat demo of real-time outputs. You should follow these instructions to manage the datasets.

  • Data Download Instructions

    • download some useful data/scripts pre-collected by us.
      • unzip benchmark_data.zip && cd benchmark_data
      • bmk_dir=${work_dir}/data/benchmark_data
    • gqa
      • download its image data following the official instructions here
      • cd ${bmk_dir}/gqa && ln -s /path/to/gqa/images images
    • mme
      • download the data following the official instructions here.
      • cd ${bmk_dir}/mme && ln -s /path/to/MME/MME_Benchmark_release_version images
    • pope
      • download coco from POPE following the official instructions here.
      • cd ${bmk_dir}/pope && ln -s /path/to/pope/coco coco && ln -s /path/to/coco/val2014 val2014
    • sqa
      • download images from the data/scienceqa folder of the ScienceQA repo.
      • cd ${bmk_dir}/sqa && ln -s /path/to/sqa/images images
    • textvqa
      • download images following the instructions here.
      • cd ${bmk_dir}/textvqa && ln -s /path/to/textvqa/train_images train_images
    • mmbench
      • no action is needed.

• organize the data directory as follows after downloading all of them:

  • Data Structure Tree

    .
    ├── benchmark_data
    │   ├── gqa
    │   │   ├── convert_gqa_for_eval.py
    │   │   ├── eval.py
    │   │   ├── images -> /path/to/your/gqa/images
    │   │   ├── llava_gqa_testdev_balanced.jsonl
    │   │   └── testdev_balanced_questions.json
    │   ├── mmbench
    │   │   ├── convert_mmbench_for_submission.py
    │   │   ├── eval.py
    │   │   └── mmbench_dev_en_20231003.tsv
    │   ├── mme
    │   │   ├── calculation.py
    │   │   ├── convert_answer_to_mme.py
    │   │   ├── images -> /path/to/your/MME/MME_Benchmark_release_version
    │   │   └── llava_mme.jsonl
    │   ├── pope
    │   │   ├── coco -> /path/to/your/pope/coco
    │   │   ├── eval.py
    │   │   ├── llava_pope_test.jsonl
    │   │   └── val2014 -> /path/to/your/coco/val2014
    │   ├── sqa
    │   │   ├── eval.py
    │   │   ├── images -> /path/to/your/scienceqa/images
    │   │   ├── llava_test_CQM-A.json
    │   │   ├── pid_splits.json
    │   │   └── problems.json
    │   └── textvqa
    │       ├── eval.py
    │       ├── llava_textvqa_val_v051_ocr.jsonl
    │       ├── TextVQA_0.5.1_val.json
    │       └── train_images -> /path/to/your/textvqa/train_images
    ├── finetune_data
    │   ├── llava_v1_5_mix665k.json
    │   ├── Beta-KD_V2_FT_Mix2M.json
    │   ├── coco
    │   │   ├── train2017
    │   │   └── val2017
    │   ├── gqa
    │   │   └── images
    │   ├── iconqa_data
    │   │   └── iconqa
    │   │       └── train
    │   │           ├── choose_img
    │   │           ├── choose_txt
    │   │           └── fill_in_blank
    │   ├── ocr_vqa
    │   │   └── images
    │   ├── sam
    │   │   └── images
    │   ├── SBU
    │   │   └── images
    │   ├── ScienceQA
    │   │   └── train
    │   ├── share_textvqa
    │   │   └── images
    │   ├── textvqa
    │   │   └── train_images
    │   ├── vg
    │   │   ├── VG_100K
    │   │   └── VG_100K_2
    │   ├── web-celebrity
    │   │   └── images
    │   ├── web-landmark
    │   │   └── images
    │   └── wikiart
    │       └── images
    └── pretrain_data
        ├── share-captioner_coco_lcs_sam_1246k_1107.json
        ├── blip_laion_cc_sbu_558k.json
        ├── images
        ├── coco
        │   └── train2017
        ├── llava
        │   └── llava_pretrain
        └── sam
            └── images
    

2. Training and Evaluation

2.1 Evaluation

Evaluate a Hugging Face model on a single dataset (e.g. ScienceQA):

bash scripts/benchmark.sh jsun39/Cosine-Beta-KD-Instance eval-results "sqa"

Evaluate a local finetuned model on a single dataset (e.g. ScienceQA):

bash scripts/benchmark.sh outputs-finetune/finetune/checkpoint-18000 eval-results "sqa"

Evaluate a Hugging Face model on all 6 benchmark datasets:

bash scripts/benchmark.sh jsun39/Cosine-Beta-KD-Instance eval-results

Evaluate a local finetuned model on all 6 benchmark datasets:

bash scripts/benchmark.sh outputs-finetune/finetune/checkpoint-18000 eval-results

2.2 Training

Training Overview

The training process of Beta-KD V2 is divided into two stages:

• stage I: pre-training

  • Setup: ❄️ frozen vision encoder + 🔥 learnable LDP V2 projector + ❄️ frozen LLM
  • Time: around 21 hours on 4x H100 (80G), or 25 hours on 4x A100 (80G)
  • Batch setting:

    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 32
    

  • Memory: approximately 29G / 80G GPU memory

• stage II: multi-task training

  • Setup: ❄️ frozen vision encoder + 🔥 learnable LDP V2 projector + 🔥 learnable LLM
  • Time: around 14 hours on 4x H100 (80G), or 27 hours on 4x A100 (80G)
  • Batch setting:

    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 32
    

  • Memory: approximately 29G / 80G GPU memory

Default setting — follows the AlignKD paper, with equal loss weighting across two training stages:

bash 1_pretrain.sh 1 align-kd equal
bash 2_finetune.sh 1 align-kd equal

Beta-KD with task-level uncertainty weighting:

bash 1_pretrain.sh 1 cosine-probs task
bash 2_finetune.sh 1 cosine-probs task

Beta-KD with instance-conditional weighting:

bash 1_pretrain.sh 1 cosine-probs instance
bash 2_finetune.sh 1 cosine-probs instance

All available argument options:

🤝 Acknowledgments

• Beta-KD: the codebase we built upon. Thanks for their wonderful work! 👏
• TAID: some kl loss implementation are borrow from this code base!

✏️ Reference

If you find Beta-KD useful in your research or applications, please consider giving a star ⭐ and citing using the following BibTeX:

@inproceedings{sun2026uncertainty,
  title={Uncertainty-Aware Knowledge Distillation for Multimodal Large Language Models},
  author={Sun, Jingchen and Han, Shaobo and Patel, Deep and Kohno, Wataru and Jin, Can and Chen, Changyou},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2026}
}