fagp / sinkhorn-rebasin

MIT License
15 stars 1 forks source link

Re-basin via implicit Sinkhorn differentiation

Implementation of paper Re-basin via implicit Sinkhorn differentiation (Accepted at CVPR 2023).

Installation

pip install sinkhorn-rebasin

Running the examples

Basics Basics
Models alignment Models alignment
Linear mode connectivity Linear mode connectivity

Models alignment

cd examples
python main_alignment_{mlp|cnn|resnet}.py
Example Layer from $\theta_A$ Layer from $\pi_{\mathcal{P}}(\theta_A)$ Layer from $\theta_B$
MLP Alignment modelA MLP Alignment pi_modelA MLP Alignment modelB MLP
VGG Alignment modelA VGG Alignment pi_modelA VGG Alignment modelB VGG
ResNet18 Alignment modelA ResNet Alignment pi_modelA ResNet Alignment modelB ResNet

Linear mode connectivity

cd examples
python main_lmc_{mlp|cnn|resnet}.py
Dataset Model Accuracy LMC Cross Entropy Loss LMC
Mnist MLP LMC MLP Accuracy LMC MLP Loss
Mnist VGG LMC VGG Accuracy LMC VGG Loss
Imagenette-320 ResNet18 LMC ResNet Accuracy LMC ResNet Loss