merge-and-rebase is a research codebase for model merging, task-vector transport, and configurable fine-tuning across vision and text models. It is built for fast iteration on checkpoint merging, rebasing, and evaluation workflows.

Code for reproducing the results of DELTA: Distilling Linearized Behavior into Non-Linear Fine-Tuning for Effective Task Arithmetic (paper) and TAK: Dataless Weight Disentanglement in Task Arithmetic via Kronecker-Factored Approximate Curvature (paper) will be available soon.

• Config-driven OpenCLIP fine-tuning across multiple datasets
• Vision checkpoint merging with weighted_average, task_arithmetic, ties_merge, dare_merge, tsv_merge, isoc_merge, isocts_merge, cart_merge, and pcb/pcb_merge
• Zero-shot and merged evaluation on full benchmark test sets
• Task-vector transport utilities under merge_and_rebase.rebase
• GradFix-based rebasing for cross-base transfer in vision models
• Text fine-tuning and merge evaluation for NLI-style LLM setups

Create an environment with uv and install the package in editable mode:

Minimal install:

Install with dataset dependencies for fine-tuning and evaluation:

Install with development and test extras:

Vision fine-tuning is driven by config files. The CLI is used to choose the config and optionally restrict datasets.

Default vision config:

src/merge_and_rebase/finetune/configs/vision.yaml

This file defines:

• model backbone and pretrained variant
• training hyperparameters
• strategy choice, such as linear_probe or full
• dataset list and suite selection
• per-dataset overrides

Run all datasets defined by the config:

Restrict to a subset of datasets:

Or run a named suite:

For linear-attention LoRA runs, use the preset config:

Outputs are saved to src/checkpoints/finetune/<model>/<pretrained>/<task>/ by default.

Each task produces:

• <strategy>.pt: fine-tuned model checkpoint
• <strategy>.json: training log with metrics and hyperparameters

Each CLI run also writes:

• a structured run summary JSON
• a sibling *.events.jsonl file with append-only metric events

Text fine-tuning follows the same common plus per-dataset override structure and supports:

• full
• linear_probe
• peft_lora
• task-wise head extraction into heads.pt
• PEFT adapter export with save_format: peft

Starter configs:

• src/merge_and_rebase/finetune/configs/text.yaml
• src/merge_and_rebase/finetune/configs/text-peft.yaml

Run:

All main CLI entrypoints support the same logging config block:

By default, local structured logs are always written. If local_log_dir is unset, run summaries go to the entrypoint's natural output area when one exists, otherwise to src/.cache/run_logs/<entrypoint>/.

All entrypoints also accept CLI overrides:

• --use-wandb / --no-use-wandb
• --wandb-project
• --wandb-entity
• --wandb-tags tag1,tag2
• --wandb-mode online|offline|disabled
• --local-log-dir
• --run-name
• --log-every-n-steps

Example enabling W&B for vision fine-tuning:

Strategies live in src/merge_and_rebase/finetune/strategies/.

Currently available:

• full: full fine-tuning of all model parameters
• linear_probe: train only a linear classification head on frozen features
• ntk: first-order linearized fine-tuning via JVP around initialization
• peft_lora: LoRA adapters on the image encoder

To add a new strategy:

1. Create a new file in src/merge_and_rebase/finetune/strategies/.
2. Implement configure(...) to return the optimizer, info dict, and scheduler.
3. Register the strategy in registry.py.

Schedulers and optimizers are owned by the strategy implementation.

Checkpoint merging and evaluation code lives in src/merge_and_rebase/merge and src/merge_and_rebase/eval.

Task-vector rebasing and transport code lives in src/merge_and_rebase/rebase.

Built-in merge methods:

• weighted_average: weighted checkpoint averaging relative to a base model
• task_arithmetic: sum of task vectors relative to a base model
• ties_merge: TIES-style sparse sign-resolved task-vector merge
• dare_merge: DARE-style random sparsification with optional rescaling
• tsv_merge: task-singular-vector merge for matrix-valued weights
• isoc_merge: isotropic composition on matrix-valued deltas
• isocts_merge: isotropic merge with common and task-specific subspaces
• cart_merge: low-rank CART merge for 2D parameters
• pcb and pcb_merge: PCB merge on flattened task vectors

For tsv_merge, isoc_merge, and isocts_merge, 2D parameters use the method-specific SVD rule. 1D parameters default to zero task-vector deltas, matching the current repo behavior. To use the paper-style average for 1D parameters instead, set:

The functional API in merge_and_rebase.merge.methods.functional mirrors these methods through merge_* helpers for direct Python use.

Supported evaluation features:

• alpha search for merged evaluation with caching
• multi-parameter search across alpha and merge method_params with sequential grids or Sobol refinement

Evaluation is always performed on the full test set of each dataset.

Typical evaluation entrypoint:

By default, vision merge checks each tuned checkpoint for tuned_text_features and uses them when present. If they are missing, it falls back to zero-shot templates.

To override this behavior:

Or force strict checkpoint-only text features:

Command-line parameters always override config values.

Multi-parameter search is configured with hyperparam_search. Sequential search evaluates ordered method-parameter combinations and sweeps alpha inside each combination. Sobol search samples a coarse low-discrepancy set, then refines around the best region.

Example:

Sobol example:

For Llama-style causal LMs on NLI tasks such as SNLI, MNLI, SICK, QNLI, RTE, and SciTail:

Starter config:

configs/llm_merge_llama3_8b_knots_hf.json

The rebase package transports a merged task vector from a source base model to a target base model.

Core API:

• merge_task_vectors(...): computes Δ_merge = Σ_i w_i (tuned_i - source_base)
• transport_task_vector(...): transports Δ_merge from source to target coordinates
• rebase_merged_task_vectors(...): end-to-end helper that merges, transports, and applies the result

Built-in transport methods:

• identity: no-op transport where Δ' = Δ
• orthogonal_shift: removes the component of Δ aligned with (target_base - source_base)
• gradfix: masks each task-vector component using gradient signs computed on the target model
• theseus: transports matrix-shaped updates with activation-aligned orthogonal Procrustes maps

Minimal Python usage:

To add a new transport rule, create a class in src/merge_and_rebase/rebase/methods/ that implements transport(...), then register it in that module.

GradFix masks each task vector before applying it to a different pretrained base in a cross-base A to B transport setting. For each parameter, the sign of its task-vector component is compared to the gradient sign obtained on B's training data. Mismatches are either zeroed out in normal mode or forced to the gradient sign in force mode.

Pipeline:

Multi-task run over all 8 datasets:

configs/vision8_gradfix_rebase.json sets "tasks": "all", includes all 8 tuned checkpoints, and sweeps alpha from 0.0 to 2.0.

Single-task example, such as Cars:

configs/vision_gradfix_single_task.json sets "tasks": "Cars" and lists only the Cars checkpoint in tuned_ckpts.

You can also select a subset of tasks:

CLI flags always override config values.

Key config fields: