⚡️👁️ RL makes MLLMs see better than SFT

March 8, 2026 · View on GitHub

PIVOT_explain_1 PIVOT_explain_2

arXiv Website

0. Contents

  1. Installation
  2. Dataset
  3. Train
  4. Eval
  5. Citation

1. Installation

# bash
git clone https://github.com/junha1125/PIVOT.git
cd PIVOT
conda create -n pivot python=3.10.12 -y
conda activate pivot
pip install --upgrade pip  # enables PEP 660 support for editable installs
pip install -e ".[train]"
pip install --upgrade "Pillow[webp]>=10.0.0"
pip install huggingface-hub==0.31.1

We highly recommend using the base Docker image nvcr.io/nvidia/pytorch:24.01-py3.
If you use this container, remove torch/torchvision entries from 'pyproject.toml' and then run 'pip install' to avoid version conflicts.

2. Dataset

StageDescriptionSource
Stage1 Pretraining projector-onlyBLIP/LAION/CC/SBU 558k manifesthuggingface link
Stage1 Pretraining full trainingLLaVA-OneVision-Datahuggingface link
Stage2 Post-trainingMMPR-v1.2huggingface link

2.2 Download and Prepare Datasets

Step 1. Download datasets using snapshot_download

Move to the project root and install the required dependency:

cd PIVOT
pip install --upgrade huggingface_hub

Then, launch a Python interpreter (python) and run:

from huggingface_hub import snapshot_download
snapshot_download(repo_id="liuhaotian/LLaVA-Pretrain", repo_type="dataset", local_dir="datasets/BLIP558K") 
snapshot_download(repo_id="lmms-lab/LLaVA-NeXT-Data", repo_type="dataset", local_dir="datasets/LLaVA-Next") 
snapshot_download(repo_id="lmms-lab/LLaVA-OneVision-Data", repo_type="dataset", local_dir="datasets/OneVisionData")
snapshot_download(repo_id="OpenGVLab/MMPR-v1.2", repo_type="dataset", local_dir="datasets/MMPR-v1.2")

Step 2. Preprocess downloaded datasets

(1) LLaVA-Pretrain (BLIP558K)

Extract the image archive:

unzip datasets/BLIP558K/images.zip -d datasets/BLIP558K/images

(2) LLaVA-NeXT

The LLaVA-NeXT dataset includes samples stored in Parquet format. Convert them into JSON format using the provided script:

python datasets/data-processor/llava_next-unparquet.py \
    datasets/LLaVA-Next \
    llava_next_fit_mix_filtered_text_wild_738590.json

(3) LLaVA-OneVision

The LLaVA-OneVision dataset contains mixed data formats and requires multiple preprocessing steps:

  • llava_ov-1-unparquet.py Converts samples stored in Parquet format.
  • llava_ov-2-unzip.py Extracts image archives stored in ZIP format.
  • llava_ov-3-add_placeholder.py Fixes samples that are missing the <image> placeholder in the prompt. Without this placeholder, the model may be trained without attending to the image (e.g., User: Where is the iPhone?User: <image> Where is the iPhone?).

Run the preprocessing scripts in the following order:

python datasets/data-processor/llava_ov-1-unparquet.py datasets/OneVisionData  
python datasets/data-processor/llava_ov-2-unzip.py datasets/OneVisionData 
python datasets/data-processor/llava_ov-3-add_placeholder.py  

(4) MMPR (MPO)

Reconstruct and extract the image and annotation archives:

cd datasets/MMPR-v1.2
cat images.zip_* > images.zip
unzip -q annotations.zip  && unzip -q images.zip -d images  
cd ../../ # back to the root, PIVOT

Generate the JSON files required for post-training:

python datasets/data-processor/mpo-make_jsons.py --seed 777 --data_size 20000

3. Train

3.1 MLLM training

The training pipeline described in Section 3 of our paper consists of two stages:

  • Stage 1 (Pre-training) builds a base MLLM by first aligning visual–language embeddings via projector-only training, then performing end-to-end training on diverse VL datasets.
  • Stage 2 (Post-training) further refines the base model with either SFT or DPO, updating all parameters including the vision encoder.

Note: Before running, verify that --data_path and --image_folder in each script match the dataset paths from Section 2.

Stage 1: Pre-training

# Projector-only pre-training
bash scripts/train/stage1-proj-train.sh

# Full end-to-end pre-training
bash scripts/train/stage1-full-train.sh

Stage 2: Post-training

# SFT post-training
bash scripts/train/stage2-sft.sh

# DPO post-training (PIVOT)
bash scripts/train/stage2-dpo.sh

3.2 ⭐️MLLM training with a PIVOT-enhanced vision encoder⭐️

As described in Section 5 of our paper, a vision encoder refined through PIVOT (Stage 2 DPO) can be transplanted back into a freshly initialized MLLM. This decouples the vision encoder improvement from the language model, allowing any LLM backbone to benefit from PIVOT-enhanced visual representations.

The procedure mirrors the two-step pre-training of Section 3.1 — projector-only training followed by full end-to-end training — but loads the vision tower weights from a previously trained MLLM checkpoint.

# Step 1: Projector-only training (vision encoder frozen, projector randomly initialized)
bash scripts/train/pivot-proj-train.sh

# Step 2: Full end-to-end training
bash scripts/train/pivot-full-train.sh

Note: Two key arguments control the vision encoder transplant:

  • --vision_source_mllm_path: Path to the source MLLM checkpoint containing the PIVOT-enhanced vision encoder (e.g., the Stage 2 DPO checkpoint).
  • --mm_projector_source_load_layers: Number of projector layers to load from the source checkpoint. Set to 0 to randomly initialize the projector (default in the provided scripts), or to a positive integer (e.g., 1) to reuse the first N linear layers.
# e.g., ./checkpoints/stage2/stage2-dpo-siglip2-so400m-Qwen2.5-1.5B-llava
--vision_source_mllm_path ./checkpoints/stage2/<your-stage2-dpo-checkpoint>
--mm_projector_source_load_layers 0

4. Eval

We provide evaluation scripts covering both MLLM benchmarks and vision encoder representation analysis. For more details, refer to the official repository: https://github.com/cambrian-mllm/cambrian/tree/main/eval.

4.1 Simple inference

You can quickly test your trained model with a simple inference script:

python archive/LLaVA_OneVision_Tutorials.py

Note: Update the pretrained variable inside the script to point to your own trained checkpoint (e.g., /home/user/PIVOT/checkpoints/stage2/stage2-sft-siglip2-so400m-Qwen2.5-1.5B-llava).

4.2 Cambrian evaluation

The Cambrian evaluation suite covers the following benchmarks:

  • General: MME, MMBench (en + cn), SEED, GQA
  • Knowledge: ScienceQA, MMMU, MathVista, AI2D
  • OCR & Chart: ChartQA, OCRBench, TextVQA, DocVQA
  • Vision-Centric: MMVP, RealWorldQA, CVBench-2D (COCO + ADE), CVBench-3D (Omni)

Setup

cd cambrian-eval
pip install mmh3
sudo apt-get install -y libwebp-dev libjpeg-dev zlib1g-dev
pip install --upgrade "Pillow[webp]>=10.0.0"
pip install nltk
pip install openpyxl

Run

# e.g., /home/user/PIVOT/checkpoints/stage2/stage2-sft-siglip2-so400m-Qwen2.5-1.5B-llava
#        /home/user/PIVOT/checkpoints/stage2/stage2-dpo-siglip2-so400m-Qwen2.5-1.5B-llava
bash scripts/run_all_benchmarks.sh <path-to-mllm-checkpoint-root>
python scripts/tabulate.py  # aggregate results into a table

4.3 lmms-eval

Setup

git clone https://github.com/EvolvingLMMs-Lab/lmms-eval.git
# Strongly recommended: replace pyproject.toml to avoid conflicts with packages already installed for LLaVA training (torch, accelerate, etc.)
cp archive/lmms-eval/pyproject.toml lmms-eval/pyproject.toml

cd lmms-eval
pip install -e .
pip install httpx==0.23.3 protobuf==3.20 langdetect immutabledict
import nltk
nltk.download('punkt')
nltk.download('punkt_tab')
nltk.download('stopwords')
nltk.download('averaged_perceptron_tagger')

Run

For detailed execution examples, see llava_onevision.sh. You can list all available tasks with python -m lmms_eval --tasks list (e.g., HallusionBench, NaturalBench, MME_RealWorld, docvqa_val). Specify the tasks you want by joining them with commas in the --tasks argument. If multiple GPUs are visible via CUDA_VISIBLE_DEVICES, set --num_processes to match the number of GPUs.

python -m accelerate.commands.launch \
  --num_processes=1 \
  --main_process_port 20006 \
  -m lmms_eval \
  --model_args pretrained=<path-to-mllm-checkpoint-root>,conv_template=qwen_1_5,model_name=llava_qwen \
  --model llava_onevision \
  --tasks docvqa_val,hallusion_bench_image \
  --batch_size 1

Note: If pip install -e ".[train]" from Section 1 (Installation) completed successfully and llava is properly registered in your Python path, the model loading in lmms_eval/models/simple/llava_onevision.py should work without any issues.

4.4 Vision encoder analysis

The vision-only experiments in our paper directly use the following open-source repositories:

To extract the vision encoder (ViT) + projector from a trained MLLM checkpoint, use extract_vision_encoder from archive/extract_vision_encoder.py:

from archive.extract_vision_encoder import extract_vision_encoder

mllm_path = "<path-to-mllm-checkpoint-root>"
vit, projector, processor = extract_vision_encoder(mllm_path, device="cuda")

# Test inference
img = Image.open("archive/simpson.jpg").convert("RGB")
pixel_values = processor.preprocess(img, return_tensors="pt")["pixel_values"].to("cuda")

with torch.no_grad():
    vit_feats = vit(pixel_values=pixel_values, output_hidden_states=True).hidden_states[-1]
    proj_feats = projector(vit_feats)

print(vit_feats.shape)   # (1, 576, 1152)
print(proj_feats.shape)  # (1, 576, 1536)

Citation

If you find this repository helpful, please cite our paper:

@inproceedings{song2025rlseebetter,
  title   = {RL makes MLLMs see better than SFT},
  author  = {Junha Song and Sangdoo Yun and Dongyoon Han and Jaegul Choo and Byeongho Heo},
  booktitle={ICLR},
  year={2026}
}