Upload PRISM model-only checkpoints

README.md

@@ -1,3 +1,66 @@
 ---
-license: apache-2.0
+license: mit
+tags:
+- computer-vision
+- multi-task-learning
+- vision-foundation-models
+- mixture-of-experts
+- icml-2026
+datasets:
+- imagenet-1k
+- pascal-context
+- nyud-v2
 ---
+
+# PRISM Checkpoints
+
+This repository hosts model-only checkpoints for **PRISM: Synergizing Vision Foundation Models via Self-organized Expert Specialization**.
+
+Paper: https://arxiv.org/abs/2606.03444
+
+Code: https://github.com/robotyingtang/PRISM-VFM
+
+## Files
+
+```text
+checkpoints/
+  prism_stage1_vit_b.pth
+  prism_stage1_vit_l.pth
+  prism_stage2_pascal_vit_b.pth
+  prism_stage2_pascal_vit_l.pth
+  prism_stage2_nyud_vit_b.pth
+configs/
+  prism_stage1_vit_b.yml
+  prism_stage1_vit_l.yml
+  prism_stage2_pascal_vit_b.yml
+  prism_stage2_pascal_vit_l.yml
+  prism_stage2_nyud_vit_b.yml
+```
+
+The `.pth` files contain only the `model` state dict. Optimizer state, scheduler state, epoch counters, and other training metadata are removed.
+
+## Usage
+
+```bash
+hf download iiirobot/PRISM-VFM checkpoints/prism_stage2_pascal_vit_b.pth --local-dir pretrain/prism
+```
+
+```bash
+python test_condition_moe.py \
+  --exp prism_s2_pascal \
+  --config_path configs/s2_prism/pascal_s2.yml \
+  --checkpoint pretrain/prism/checkpoints/prism_stage2_pascal_vit_b.pth \
+  --results_dir results \
+  --evaluate
+```
+
+## Citation
+
+```bibtex
+@inproceedings{tang2026prism,
+  title={PRISM: Synergizing Vision Foundation Models via Self-organized Expert Specialization},
+  author={Ying Tang and Dong Li and Youjia Zhang and Zikai Song and Junqing Yu and Wei Yang},
+  booktitle={Proceedings of the 43rd International Conference on Machine Learning},
+  year={2026}
+}
+```

checkpoints/prism_stage1_vit_b.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e815c3d1963b67976b60397fde0226c9ba05707f537325bc7a8ab751c74659a
+size 1648402728

checkpoints/prism_stage1_vit_l.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:258e9d371686ccce93768233bbfe34da7830f862845924c54ade0e21e9b340f5
+size 3515493368

checkpoints/prism_stage2_nyud_vit_b.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dafbb49b60518e82e7b5c692d0fcb6c09446146a1c3073fc0888e7976d363e0a
+size 1744363606

checkpoints/prism_stage2_pascal_vit_b.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5f5d7733e33870ca7d6dea2ef778b744a7f9c36508285f9f357aea71e9c541d
+size 1768000682

checkpoints/prism_stage2_pascal_vit_l.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:563220d9a87b6e105444339d5e92c6bffc3e61868ba486d51a2c44735d3160dc
+size 3727375722

configs/prism_stage1_vit_b.yml

@@ -0,0 +1,62 @@
+results_dir: ./results/
+
+dataset: imagenet
+tr_batch: 20
+val_batch: 20
+nworkers: 4
+
+teachers:
+  dinov2:
+    backbone_type: dinov2_large
+    checkpoint_path: pretrain/vfm/dinov2_vitl14_pretrain.bin
+  clip:
+    backbone_type: timmclip_large
+    checkpoint_path: pretrain/vfm/vit_large_patch16_clip_336.bin
+  sam:
+    backbone_type: sam_large
+    checkpoint_path: pretrain/vfm/samvit_large_patch16.bin
+student:
+  backbone:
+    backbone_type: condition_moe_prism
+    vit_name: vit_base
+    vit_pretrained: False
+    down_ratio: 4
+    aligner: True
+    freeze_vit: False
+
+    moe_type: gated
+    router_type: no_structure_router
+
+    gate_constraint_ranges: [ [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ] ]
+    condition_dim:
+    num_moe_experts: 15
+    moe_top_k: 3
+    expert_hidden_ratio: 4.0
+
+    moe_layers_indices: [2, 5, 8, 11]
+    noisy_gating: True
+    vfm_projection_configs:
+    vfm_projector_type: mlp
+    vit_checkpoint_path: 
+
+    out_indices_cfg_for_task:
+      small: [ 2, 5, 8, 11 ]
+      base:  [2, 5, 8, 11]
+      large: [ 5, 11, 17, 23 ]
+      huge: [ 7, 15, 23, 31 ]
+      giant: [ 9, 19, 29, 39 ]
+      so: [ 6, 13, 20, 26 ]
+
+    vit_type: conditioned_moe
+teacher_output_indices:  [5, 11, 17, 23]
+loss_type: cos+l1+decorrelation
+
+max_epochs: 30
+base_lr: 1e-3
+weight_decay: 0.01
+warmup_epochs: 2
+grad_clip: 1.0
+save_freq: 5

configs/prism_stage1_vit_l.yml

@@ -0,0 +1,62 @@
+results_dir: ./results/
+
+dataset: imagenet
+tr_batch: 64
+val_batch: 64
+nworkers: 8
+
+teachers:
+  dinov2:
+    backbone_type: dinov2_large
+    checkpoint_path: pretrain/vfm/dinov2_vitl14_pretrain.bin
+  clip:
+    backbone_type: timmclip_large
+    checkpoint_path: pretrain/vfm/vit_large_patch16_clip_336.bin
+  sam:
+    backbone_type: sam_large
+    checkpoint_path: pretrain/vfm/samvit_large_patch16.bin
+student:
+  backbone:
+    backbone_type: condition_moe_prism
+    vit_name: vit_large
+    vit_pretrained: False
+    down_ratio: 4
+    aligner: True
+    freeze_vit: False
+
+    moe_type: gated
+    router_type: no_structure_router
+
+    gate_constraint_ranges: [ [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ] ]
+    condition_dim:
+    num_moe_experts: 15
+    moe_top_k: 3
+    expert_hidden_ratio: 4.0
+
+    moe_layers_indices: [ 5, 11, 17, 23 ]
+    noisy_gating: True
+    vfm_projection_configs:
+    vfm_projector_type: mlp
+    vit_checkpoint_path: 
+
+    out_indices_cfg_for_task:
+      small: [ 2, 5, 8, 11 ]
+      base:  [2, 5, 8, 11]
+      large: [ 5, 11, 17, 23 ]
+      huge: [ 7, 15, 23, 31 ]
+      giant: [ 9, 19, 29, 39 ]
+      so: [ 6, 13, 20, 26 ]
+
+    vit_type: conditioned_moe
+teacher_output_indices:  [5, 11, 17, 23]
+loss_type: cos+l1+decorrelation
+
+max_epochs: 30
+base_lr: 1e-3
+weight_decay: 0.01
+warmup_epochs: 2
+grad_clip: 1.0
+save_freq: 5

configs/prism_stage2_nyud_vit_b.yml

@@ -0,0 +1,67 @@
+results_dir: ./results/
+
+dataset: nyud
+tr_batch: 4
+val_batch: 8
+nworkers: 4
+task_dict:
+  semseg: 1
+  normals: 1
+  edge: 1
+  depth: 1
+
+teachers:
+
+  dinov2:
+    backbone_type: dinov2_large
+    checkpoint_path: pretrain/vfm/dinov2_vitl14_pretrain.bin
+  clip:
+    backbone_type: timmclip_large
+    checkpoint_path: pretrain/vfm/vit_large_patch16_clip_336.bin
+  sam:
+    backbone_type: sam_large
+    checkpoint_path: pretrain/vfm/samvit_large_patch16.bin
+
+student:
+  backbone:
+    backbone_type: condition_moe_prism
+    vit_name: vit_base
+    vit_pretrained: False
+    down_ratio: 4
+    aligner: False
+    freeze_vit: False
+    moe_type: gated
+    router_type: no_structure_router
+    gate_constraint_ranges: [ [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ] ]
+
+    condition_dim:
+    num_moe_experts: 15
+    expert_hidden_ratio: 4.0
+    moe_top_k: 3
+    moe_layers_indices: [ 2, 5, 8, 11 ]
+    noisy_gating: True
+    vfm_projection_configs:
+    vfm_projector_type: mlp
+    vit_checkpoint_path:
+    vfm_training: True
+    out_indices_cfg_for_task:
+      small: [ 2, 5, 8, 11 ]
+      base: [ 2, 5, 8, 11 ]
+      large: [ 5, 11, 17, 23 ]
+      huge: [ 7, 15, 23, 31 ]
+      giant: [ 9, 19, 29, 39 ]
+      so: [ 6, 13, 20, 26 ]
+    vit_type: conditioned_moe
+
+teacher_output_indices:  [5, 11, 17, 23]
+loss_type: cos+l1+decorrelation
+
+max_epochs: 0
+max_iters: 40000
+eval_freq: 2
+base_lr: 2e-5
+weight_decay: 1e-6
+grad_clip: 10

configs/prism_stage2_pascal_vit_b.yml

@@ -0,0 +1,71 @@
+results_dir: ./results/
+
+dataset: pascalcontext
+tr_batch:  3
+val_batch: 8
+nworkers: 2
+task_dict:
+  semseg: 1
+  human_parts: 1
+  normals: 1
+  edge: 1
+  sal: 1
+
+teachers:
+  dinov2:
+    backbone_type: dinov2_large
+    checkpoint_path: pretrain/vfm/dinov2_vitl14_pretrain.bin
+    embed_dim: 1024
+  clip:
+    backbone_type: timmclip_large
+    checkpoint_path: pretrain/vfm/vit_large_patch16_clip_336.bin
+    embed_dim: 1024
+  sam:
+    backbone_type: sam_large
+    checkpoint_path: pretrain/vfm/samvit_large_patch16.bin
+    embed_dim: 1024
+
+student:
+  backbone:
+    backbone_type: condition_moe_prism
+    vit_name: vit_base
+    vit_pretrained: False
+    down_ratio: 4
+    aligner: False
+    freeze_vit: False
+
+    moe_type: gated
+    router_type: no_structure_router
+    gate_constraint_ranges: [ [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ] ]
+
+    condition_dim:
+    num_moe_experts: 15
+    expert_hidden_ratio: 4.0
+    moe_top_k: 3
+    moe_layers_indices: [2, 5, 8, 11]
+    noisy_gating: True
+    vfm_projection_configs:
+    vfm_projector_type: mlp
+    vit_checkpoint_path:
+    vfm_training: True
+    out_indices_cfg_for_task:
+      small: [ 2, 5, 8, 11 ]
+      base:  [2, 5, 8, 11]
+      large: [ 5, 11, 17, 23 ]
+      huge: [ 7, 15, 23, 31 ]
+      giant: [ 9, 19, 29, 39 ]
+      so: [ 6, 13, 20, 26 ]
+    vit_type: conditioned_moe
+
+teacher_output_indices:  [5, 11, 17, 23]
+loss_type: cos+l1+decorrelation
+
+max_epochs: 50
+max_iters: 40000
+eval_freq: 2
+base_lr: 2e-5
+weight_decay: 1e-6
+grad_clip: 10

configs/prism_stage2_pascal_vit_l.yml

@@ -0,0 +1,68 @@
+results_dir: ./results/
+
+dataset: pascalcontext
+tr_batch: 6
+val_batch: 12
+nworkers: 4
+task_dict:
+  semseg: 1
+  human_parts: 1
+  normals: 1
+  edge: 1
+  sal: 1
+
+teachers:
+  dinov2:
+    backbone_type: dinov2_large
+    checkpoint_path: pretrain/vfm/dinov2_vitl14_pretrain.bin
+  clip:
+    backbone_type: timmclip_large
+    checkpoint_path: pretrain/vfm/vit_large_patch16_clip_336.bin
+  sam:
+    backbone_type: sam_large
+    checkpoint_path: pretrain/vfm/samvit_large_patch16.bin
+
+student:
+  backbone:
+    backbone_type: condition_moe_prism
+    vit_name: vit_large
+    vit_pretrained: False
+    down_ratio: 4
+    aligner: False
+    freeze_vit: False
+
+    moe_type: gated
+    router_type: no_structure_router
+    gate_constraint_ranges: [ [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ],
+                              [ 0.5, 1.0 ] ]
+
+    condition_dim:
+    num_moe_experts: 15
+    expert_hidden_ratio: 4.0
+    moe_top_k: 3
+    moe_layers_indices: [ 5, 11, 17, 23 ]
+    noisy_gating: True
+    vfm_projection_configs:
+    vfm_projector_type: mlp
+    vit_checkpoint_path:
+    vfm_training: True
+    out_indices_cfg_for_task:
+      small: [ 2, 5, 8, 11 ]
+      base:  [2, 5, 8, 11]
+      large: [ 5, 11, 17, 23 ]
+      huge: [ 7, 15, 23, 31 ]
+      giant: [ 9, 19, 29, 39 ]
+      so: [ 6, 13, 20, 26 ]
+    vit_type: conditioned_moe
+
+teacher_output_indices:  [5, 11, 17, 23]
+loss_type: cos+l1+decorrelation
+
+max_epochs: 0
+max_iters: 50000
+eval_freq: 2
+base_lr: 2e-5
+weight_decay: 1e-6
+grad_clip: 10