Package: rebasin
rebasinAn implementation of methods described in"Git Re-basin"-paper by Ainsworth et al.Can be applied toarbitrary models, without modification.(Well,almostarbitrary models, seeLimitations).Table of ContentsInstallationUsageLimitationsResultsAcknowledgementsInstallationRequirements should be automatically installed, but one of them is graphviz,
which you might have to install per apt / brew / ... on your device.The following install instructions are taken directly fromtorchview's installation instructions.Debian-based Linux distro (e.g. Ubuntu):apt-getinstallgraphvizWindows:chocoinstallgraphvizmacOSbrewinstallgraphvizsee more detailshere.Then, install rebasin via pip:pipinstallrebasinUsageCurrently, only weight-matching is implemented as a method for rebasing,
and only a simplified form of linear interpolation is implemented.The following is a minimal example. For now, the documentation lives in the docstrings,
though I intend to create a proper one.PermutationCoordinateDescentandinterpolation.LerpSimpleare the main classes, besideMergeMany(see below).fromrebasinimportPermutationCoordinateDescentfromrebasinimportinterpolationmodel_a,model_b,train_dl=...input_data=next(iter(train_dl))[0]# Rebasinpcd=PermutationCoordinateDescent(model_a,model_b,input_data)# weight-matchingpcd.rebasin()# Rebasin model_b towards model_a. Automatically updates model_b# Interpolatelerp=interpolation.LerpSimple(models=[model_a,model_b],devices=["cuda:0","cuda:1"],# Optional, defaults to cpudevice_interp="cuda:2",# Optional, defaults to cpusavedir="/path/to/save/interpolation"# Optional, save all interpolated models)lerp.interpolate(steps=99)# Interpolate 99 models between model_a and model_bTheMergeMany-algorithm is also implemented
(though there will be interface-changes regarding the devices in the future):fromrebasinimportMergeManyfromtorchimportnnclassExampleModel(nn.Module):...model_a,model_b,model_c=ExampleModel(),ExampleModel(),ExampleModel()train_dl=...# Mergemerge=MergeMany(models=[model_a,model_b,model_c],working_model=ExampleModel(),input_data=next(iter(train_dl))[0],)merged_model=merge.run()# The merged model is also accessible through merge.working_model,#   but only after merge.run() has been called.TerminologyIn this document, I will use the following terminology:To rebasin: To apply one of the methods described in the paper to a model,
permuting the rows and columns of its weights (and biases)model_a: The model that stays unchangedmodel_b: The model that is changed by rebasin it towardsmodel_amodel_b (original)for the unchanged, originalmodel_bmodel_b (rebasin)for the changed, rebasinedmodel_bPath: A sequence of modules in a modelLimitationsOnly some methods are implementedFor rebasin, only weight-matching is implemented viarebasin.PermutationCoordinateDescent.For interpolation, only a simplified method of linear interpolation is implemented
viarebasin.interpolation.LerpSimple.Limitations of thePermutationCoordinateDescent-classThePermutationCoordinateDescent-class only permutes some Modules.
Most modules should work, but others may behave unexpectedly. In this case,
you need to add the module torebasin/modules.py;
make sure it is included in theinitialize_module-function
(preferably by putting it into theSPECIAL_MODULES-dict).Additionally, thePermutationCoordinateDescent-class only works withnn.Modules, not functions. There is a requirement to have the permuted model
produce the same output as the un-permutedModule, which is a pretty
tight constraint. In some models, it isn't a problem at all, but especially in
models with lots of short residual blocks, it may (but doesn't have to) be a problem.
Where it is a problem, few to no parameters get permuted, which defeats the purpose of rebasin.For example, @tysam-code'shlb-gpt, a small but fast
language model implementation, isn't permuted at all.
Vision transformers liketorchvision.models.vit_b_16have only very few permutations
applied to them. In general,transformer models don't work well, because they
reshape the input-tensor, and directly follow that up with residual blocks.
This means that almost nothing of the model can be permuted
(a single Linear layer between the reshaping and the first residual block would fix that,
but this isn't usually done...).On the other hand,CNNs usually work very well.If you are unsure, you can always print the model-graph! To do so, write:fromrebasinimportPermutationCoordinateDescentpcd=PermutationCoordinateDescent(...)print(pcd.pinit.model_graph)# pinit stands for "PermutationInitialization"ResultsFor the full results, seerebasin-results(I don't want to upload a bunch of images to this repo, so the results are in their own repo).The clearest results were produces onhlb-CIFAR10.
For results on that model, seehere.Here is a little taste of the results for that model:WhilePermutationCoordinateDescentdoesn't fully eliminate the loss-barrier,
it does reduce it significantly, and, surprisingly, even moreso for the accuracy-barrier.You can also find results for theMergeMany-algorithm there.AcknowledgementsGit Re-Basin:Ainsworth, Samuel K., Jonathan Hayase, and Siddhartha Srinivasa. 
"Git re-basin: Merging models modulo permutation symmetries." 
arXiv preprint arXiv:2209.04836 (2022).Link:https://arxiv.org/abs/2209.04836(accessed on April 9th, 2023)ImageNet:I've used the ImageNet Data from the 2012 ILSVRC competition to evaluate
the algorithms from rebasin on thetorchvision.models.Olga Russakovsky*, Jia Deng*, Hao Su, Jonathan Krause, Sanjeev Satheesh, 
Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, 
Alexander C. Berg and Li Fei-Fei. (* = equal contribution) 
ImageNet Large Scale Visual Recognition Challenge. arXiv:1409.0575, 2014Paper (link)(Accessed on April 12th, 2023)Torchvision modelsFor testing, I've used the torchvision models (v.015), of course (or I will):https://pytorch.org/vision/0.15/models.htmlHLB-CIFAR10For testing, I forkedhlb-CIFAR10by@tysam-code:authors:
- family-names: "Balsam"
  given-names: "Tysam&"
title: "hlb-CIFAR10"
version: 0.4.0
date-released: 2023-02-12
url: "https://github.com/tysam-code/hlb-CIFAR10"HLB-GPTFor testing, I also usedhlb-gptby @tysam-code:authors:
  - family-names: "Balsam"
    given-names: "Tysam&"
title: "hlb-gpt"
version: 0.0.0
date-released: 2023-03-05
url: "https://github.com/tysam-code/hlb-gpt"OtherMy code took inspiration from the following sources:https://github.com/themrzmaster/git-re-basin-pytorchI used the amazing librarytorchviewto visualize the models:https://github.com/mert-kurttutan/torchview