This repository provides the official PyTorch implementation of our T³ Medical Model-Merging paper:
T³ Workflow
Figure 1: Dynamic test-time merging workflow of T³
Official implementation of T³: Test-Time Model Merging in Vision-Language Models for Zero-Shot Medical Imaging, a method for adaptive fusion of pretrained and fine-tuned vision-language models at test time using Jensen-Shannon divergence.
Key Features
🧠 Mutual Information Guidance: Uses JS divergence to measure model consensus.
⚡ Backpropagation-Free: No gradient updates required during inference.
🏥 Medical Modality Agnostic: Validated consistency on 4x medical imaging domains.
🚀 Batch-Wise Efficiency: Reduces compute cost by 32x vs sample-wise merging.
📈 SOTA Performance: Outperforms 8+ baselines in accuracy & robustness.
JS = Jensen-Shannon divergence between pretrained and fine-tuned model predictions.
σ = Sigmoid function for smooth scaling.
γ = Scaling factor (default=0.5).
Visual Explanation of the Method
Below justifies the method and its effectiveness:
Dynamic Weighting Based on Model Agreement
We propose using Jensen–Shannon (JS) divergence to measure mutual information between pretrained (p_pt) and fine-tuned (p_ft) model predictions, offering a more robust gauge of joint confidence than entropy-based methods like DaWin's entropy ratio:
Empirically, (I(x)) correlates positively with (R(x)), but better distinguishes disagreements, validating its use for adaptive merging.
Mutual Information vs. Entropy
MI vs Entropy
Figure 3: Relationship between mutual information and entropy for adaptive merging.
Performance Across Modalities
Performance Comparison
Figure 4: T³ achieves superior performance across multiple medical imaging modalities.
Folder Structure
Do check our HuggingFace page for Expert Models and Evaluation Datasets.
T3/
├── clip/ # CLIP model adaptations
├── data/ # Data Utilities
├── utils/ # Helper functions
├── models/ # Put your finetuned models HERE
├── dataset/ # Put your medimeta/medmnist-c eval data HERE
├── baselines.py # Comparison methods
├── t_cube.py # Core T³ implementation
├── BetaMixture.py # Auxiliary models
└── README.md # This document
Reproducing Results
To reproduce the results from the paper, you can run the t_cube.py script. This script handles the evaluation of T³ and its baselines across multiple datasets and severity levels. Additional baselines are available in baselines.py.
To understand the script better; in t_cube.py:
Refer to the compute_samplewise_tcube_weights and compute_samplewise_tcube_weights_MI functions for entropy (DaWiN baseline) and Our mutual information-based merging.
Check the evaluate_on_test_set function for how datasets and severities are processed.
Explore the evaluate_tcube function for the merging and evaluation logic.
Pretrained Weights
We provide pretrained weights for the following models:
Generalist CLIP: A pretrained model for general vision-language tasks.
Expert CLIPs: 4x Fine-tuned models for the following medical imaging domains:
We provided Breast Imaging evaluation sets on HuggingFace page. Please download from there.
If you need to run multiple modalities datasets, just pass --testset arg with 'bloodmnist/breastmnist/'. This will evaluate medmnist-c and medimeta from each modality, resulting in 4 datasets evaluation.
If you need all modality datasets, you can find them as follows:
This project is licensed under the MIT License - see the LICENSE file for details.
Citation
If you find this work useful, please cite the arXiv version below:
@misc{imam2025t3testtimemodelmerging,
title={T3: Test-Time Model Merging in VLMs for Zero-Shot Medical Imaging Analysis},
author={Raza Imam and Hu Wang and Dwarikanath Mahapatra and Mohammad Yaqub},
year={2025},
eprint={2510.27265},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2510.27265},
}
Contact
For questions or collaborations, contact Raza Imam. Please feel free to raise an issue in facing error in reproducing the results.
