Skip to content

Jimmy005452/MMCR

BranchesTags

Repository files navigation

MMCR: Sequential Layer-wise Model Merging via Reinforcement Learning

Project Page Report Poster

MMCR learns sequential non-negative layer-wise coefficients for task-vector based Vision Transformer model merging. Using TIES-selected task vectors, MMCR optimizes the merge policy with reinforcement learning and achieves the highest average accuracy in the eight-dataset CLIP ViT-L/14 merging benchmark.

Main Results

Method SUN397 Cars RESISC45 EuroSAT SVHN GTSRB MNIST DTD Avg. Time (min)
TIES 76.11 83.93 83.65 80.66 84.70 89.14 89.29 65.53 81.63 3.23
AdaMerging-300 78.20 85.56 87.67 83.50 68.45 86.81 61.74 70.96 77.86 185.00
AdaMerging-500 78.71 87.58 90.21 88.14 88.13 89.37 91.45 72.87 85.81 356.87
AdaMerging-600 78.70 87.66 90.62 88.78 89.08 89.79 92.61 73.67 86.36 418.90
MMCR (ours) 77.57 89.11 89.06 85.72 90.65 89.96 92.54 76.64 86.41 245.17

The checkpoint corresponding to the 86.41 average accuracy result is available on Hugging Face: RL-final-MMCR-best

Overview

MMCR formulates layer-wise task-vector merging as a sequential coefficient-learning problem. Given TIES-selected task vectors, the policy predicts non-negative coefficients for each transformer layer and optimizes the merged encoder with an entropy-based reward.

The full workflow is:

  1. Prepare datasets: download the five torchvision-compatible datasets automatically, then download the prepared SUN397, Stanford Cars, and RESISC45 mirror.
  2. Prepare checkpoints: create the zeroshot encoder, convert pretrained source encoders, and download or generate dataset heads.
  3. Evaluate the source models.
  4. Run model-merging baselines: Task Arithmetic, TIES, DARE, NAN, and AdaMerging.
  5. Run RL-based merging: GRPO/RLOO, PPO-GAE, GSPO-like, Dr.GRPO-like, and DAPO-lite.
  6. Evaluate any merged encoder with eval.py.
  7. Optionally transfer a trained GRPO policy to another dataset combination with the same number of tasks.

Environment Setup

  • Python 3.11+
  • GPU recommended (CUDA supported)

The main experiments were developed and tested on two NVIDIA RTX 3090 GPUs.

For running this project on your local machine, follow the steps below.

Step 0: Clone the Repository

git clone https://github.com/Jimmy005452/MMCR.git
cd MMCR

Step1: Create a Conda Environment

Conda is recommended for dependency isolation.

# Create a conda environment
conda create -n mmcr python=3.11 -y

# Activate environment
conda activate mmcr

Step2: Install PyTorch

Please install PyTorch based on your system configuration:

  1. Visit: https://pytorch.org/get-started/locally/
  2. Select your OS, package (pip), and CUDA version
  3. Run the generated command

Example (CUDA 12.6):

pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu126

Step3: Install Other Dependencies

pip install -r requirements.txt

Repository Layout

MMCR/
├─ README.md
├─ CLI.md                         # runnable command examples for experiments
├─ requirements.txt
├─ eval.py                         # evaluate any merged encoder
├─ mmcr/
│  ├─ data.py                      # dataset loaders and transforms
│  ├─ models.py                    # CLIP/timm model builders
│  ├─ checkpoints.py               # checkpoint loading and saving helpers
│  ├─ task_vectors.py              # task-vector construction utilities
│  └─ evaluation.py                # shared evaluation utilities
├─ model_merging/
│  ├─ task_arithmetic.py           # Task Arithmetic baseline
│  ├─ ties_merge.py                # streaming TIES baseline
│  ├─ dare_merge.py                # DARE baseline
│  ├─ nan_merge.py                 # NAN baseline
│  └─ adamerging_modes.py          # AdaMerging task/tensor modes
├─ rl_methods/
│  ├─ source_baselines.py          # evaluate individual source models
│  ├─ rl_mmcr_GRPO_RLOO/           # GRPO/RLOO and GRPO-style variants
│  └─ rl_mmcr_PPO_GAE_Actor-Critic/ # PPO-GAE actor-critic baseline
├─ analysis_tools/
│  └─ task-vector analysis tools
├─ hf_tools/
│  └─ convert_hf_clip_to_timm.py   # convert HF CLIP checkpoints to encoder.pt
└─ source_models/
   ├─ save_zeroshot.py             # save the shared zeroshot encoder
   └─ train_vit_l14.py             # train/download-compatible dataset heads

Supported Datasets

Use these dataset keys in all commands:

Key Dataset Classes Notes
sun397 SUN397 397 Expected under data/sun397/train and data/sun397/test.
stanford_cars Stanford Cars 196 Alias is normalized internally to cars.
resisc45 NWPU-RESISC45 45 Uses split files under data/resisc45/.
eurosat EuroSAT 10 Torchvision-compatible.
svhn SVHN 10 Torchvision-compatible.
gtsrb GTSRB 43 Torchvision-compatible.
mnist MNIST 10 Converted to RGB by the loader.
dtd DTD 47 Torchvision-compatible.

For convenience, the eight-dataset order used in the experiments is:

DATASETS_8="sun397 stanford_cars resisc45 eurosat svhn gtsrb mnist dtd"

Data Preparation

This repository does not store raw datasets. Five datasets are torchvision-compatible and can be downloaded automatically by the loaders. Three datasets need the prepared directory layout used by this project.

Download Torchvision Datasets

The following datasets can be downloaded automatically when you run training or evaluation without --no-download:

eurosat
svhn
gtsrb
mnist
dtd

Download the five torchvision-compatible datasets before checkpoint-dependent commands:

python - <<'PY'
from pathlib import Path
from mmcr.data import build_transforms, make_dataset

transform = build_transforms(train=False)
for dataset in ["eurosat", "svhn", "gtsrb", "mnist", "dtd"]:
    for train in [True, False]:
        split = "train" if train else "test"
        print(f"Downloading {dataset} {split}...")
        make_dataset(dataset, Path("data"), train=train, transform=transform, download=True)
PY

After the data is present locally, pass --no-download to avoid repeated download checks.

Download Prepared Datasets from Hugging Face

SUN397, Stanford Cars, and RESISC45 require the prepared layout below. Download those three datasets into data/ from the project mirror:

export HF_XET_HIGH_PERFORMANCE=1
hf download jimmy0214/RL-final-MMCR-dataset \
  --repo-type dataset \
  --local-dir data \
  --max-workers 16

Expected Data Layout

data/
├─ sun397/
│  ├─ train/
│  │  └─ <class_name>/*.jpg
│  └─ test/
│     └─ <class_name>/*.jpg
├─ stanford_cars/
│  ├─ devkit/
│  │  ├─ cars_train_annos.mat
│  │  ├─ cars_test_annos_withlabels.mat
│  │  └─ cars_meta.mat
│  ├─ cars_train/*.jpg
│  └─ cars_test/*.jpg
├─ resisc45/
│  ├─ resisc45-train.txt
│  ├─ resisc45-val.txt
│  ├─ resisc45-test.txt
│  └─ NWPU-RESISC45/
│     └─ <class_name>/*.jpg
├─ eurosat/
├─ svhn/
├─ gtsrb/
├─ mnist/
└─ dtd/

When the data mirror is already present, pass --no-download to avoid torchvision downloads.

Checkpoint Preparation

Create Zeroshot Encoder

Create the base CLIP ViT-L/14 encoder used as the zeroshot initialization:

python -m source_models.save_zeroshot --output checkpoints/zeroshot.pt --overwrite

Download and Convert Pretrained Source Encoders

The eight pretrained Hugging Face CLIP source encoders can be downloaded and converted to this repository's timm encoder.pt format with:

python -m hf_tools.convert_hf_clip_to_timm \
  --download-root checkpoints/hf_raw \
  --output-root checkpoints \
  --overwrite

By default this converts the following repositories:

tanganke/clip-vit-large-patch14_sun397
tanganke/clip-vit-large-patch14_stanford-cars
tanganke/clip-vit-large-patch14_resisc45
tanganke/clip-vit-large-patch14_eurosat
tanganke/clip-vit-large-patch14_svhn
tanganke/clip-vit-large-patch14_gtsrb
tanganke/clip-vit-large-patch14_mnist
tanganke/clip-vit-large-patch14_dtd

Download Source Heads

We recommend directly using the heads we have trained. Please download the provided head.pt files from Hugging Face:

hf download jimmy0214/RL-final-MMCR-heads \
  --repo-type model \
  --local-dir checkpoints

We also provide the code for training the heads. For more details, please refer to the Generate Source Heads by Fine-tuning section in CLI.md.

Expected Checkpoint Layout

The checkpoint directory should follow a folder-based structure. Each dataset-specific checkpoint must be placed under its corresponding dataset folder. Every dataset folder should contain two files: encoder.pt and head.pt.

The encoder.pt file stores the task-specific encoder weights, while head.pt stores the corresponding classification head. The shared zero-shot checkpoint should be placed directly under the root checkpoints/ directory as zeroshot.pt.

checkpoints/
├─ zeroshot.pt
├─ sun397/
│  ├─ encoder.pt
│  └─ head.pt
├─ stanford_cars/
├─ resisc45/
├─ eurosat/
├─ svhn/
├─ gtsrb/
├─ mnist/
└─ dtd/

All other dataset folders follow the same structure as sun397/.

Evaluate Source Models

Evaluate each source model on its own test split and cache the source baseline accuracies:

python -m rl_methods.source_baselines \
  --datasets sun397 stanford_cars resisc45 eurosat svhn gtsrb mnist dtd \
  --checkpoint-root checkpoints \
  --data-root data \
  --batch-size 64 \
  --gpu 0 \
  --amp \
  --no-download \
  --output-json results/source_baselines_8datasets_test.json

Model Merging Baselines

For the commands used to run model merging baselines, please refer to the Model Merging Baselines section in CLI.md.

We provide instructions for running the following five baseline methods:

  • Task Arithmetic
  • TIES
  • DARE
  • NAN
  • AdaMerging (Tensor-wise)

RL-based Model Merging (GRPO on 8 Datasets)

Below are our recommended GRPO settings:

python -m rl_methods.rl_mmcr_GRPO_RLOO.train \
  --datasets sun397 stanford_cars resisc45 eurosat svhn gtsrb mnist dtd \
  --checkpoint-root checkpoints \
  --data-root data \
  --reward-mode entropy \
  --reward-split val \
  --reward-sampling-mode stratified_pool \
  --reward-pool-size 1024 \
  --selection-reward-pool-position 0 \
  --selection-candidates 1 \
  --reward-batch-size 64 \
  --reward-batches-per-dataset 1 \
  --merge-granularity layer \
  --state-mode full_coefficients \
  --coefficient-init 0.3 \
  --task-vector-mode ties \
  --top-k-percent 20 \
  --policy-hidden-dim 128 \
  --iterations 300 \
  --group-size 8 \
  --grpo-epochs 4 \
  --lr 3e-4 \
  --clip-eps 0.2 \
  --entropy-coef 0.005 \
  --target-kl 0.03 \
  --advantage-mode rloo \
  --terminal-bonus 1.0 \
  --reward-scale 1.0 \
  --step-reward-coef 0.25 \
  --score-imbalance-coef 0.5 \
  --reward-eval-interval 13 \
  --log-every 20 \
  --gpu 0 \
  --num-workers 4 \
  --seed 2033 \
  --amp \
  --no-download \
  --cache-task-vectors-device \
  --batched-reward-eval \
  --batched-reward-max-samples 128 \
  --trajectory-devices 0 1 \
  --output-dir rl_mmcr_GRPO_RLOO_runs/layer8_best_interval13_seed2033

Note: GPU Configuration for GRPO Training

--gpu specifies the main GPU used for policy training, while --trajectory-devices specifies the GPU devices used for trajectory generation.

For single-GPU execution, use the same device for both arguments:

--gpu 0 --trajectory-devices 0

For two-GPU trajectory generation, keep --gpu as the main training device and list both devices in --trajectory-devices:

--gpu 0 --trajectory-devices 0 1

GRPO writes the following artifacts to --output-dir:

encoder.pt          # exported merged encoder, controlled by --export-policy
best_encoder.pt     # merged encoder from the best selected sampled policy
final_encoder.pt    # merged encoder from the final deterministic policy
best_policy.pt      # policy checkpoint from the iteration that produced the best selected sample
final_policy.pt     # final policy checkpoint
results.json        # config, coefficients, per-iteration history, final eval, and artifact paths
training_curves.png # score/loss/KL/entropy curves
reward_curves.png   # per-episode reward curve

For experiments with other settings, please refer to the RL-based Model Merging section in CLI.md. It includes commands for the following settings:

  • GRPO on Different Numbers of Datasets
  • GRPO with Raw Task Vectors (TIES is used by default)
  • GRPO Reward Interval Ablation
  • GSPO-like
  • Dr.GRPO-like
  • DAPO-lite
  • PPO

Policy Transfer

Policy transfer applies a saved GRPO policy network to a new dataset combination with the same number of tasks. For example, a policy trained on a 2-task merge can be transferred to another 2-task merge, while a policy trained on a 3-task merge can be transferred to another 3-task merge.

The full workflow includes three steps:

  1. Train a GRPO policy on a source dataset combination.
  2. Apply the saved policy checkpoint to a new target dataset combination.
  3. Evaluate the transferred merged model.

For detailed commands and examples, please refer to the Policy Transfer section in CLI.md.

Evaluate Any Merged Encoder

Use eval.py for any encoder checkpoint, including outputs from model merging baselines, GRPO/PPO-based methods, or downloaded merged encoders.

python eval.py \
  --encoder <path/to/encoder.pt> \
  --datasets sun397 stanford_cars resisc45 eurosat svhn gtsrb mnist dtd \
  --checkpoint-root checkpoints \
  --data-root data \
  --batch-size 64 \
  --gpu 0 \
  --amp \
  --no-download \
  --results-json results/eval_custom_encoder.json \
  --results-txt results/eval_custom_encoder.txt

For additional evaluation options, including automatic dataset inference and comparison against cached source baselines, please refer to the Evaluate Any Merged Encoder section in CLI.md.

About

Sequential layer-wise model merging for Vision Transformers using reinforcement learning.

fangfangirl.github.io/RL_final/

Topics

reinforcement-learning computer-vision pytorch multi-task-learning vision-transformer model-merging grpo ties-merging task-vectors

Resources

Readme
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Contributors

Languages