fusion-bench 0.2.25__py3-none-any.whl → 0.2.27__py3-none-any.whl

fusion_bench/method/classification/image_classification_finetune.py

@@ -1,3 +1,9 @@
+"""Image Classification Fine-tuning Module.
+
+This module provides algorithms for fine-tuning and evaluating image classification models
+using PyTorch Lightning.
+"""
+
 import os
 from typing import Optional
 
@@ -23,33 +29,82 @@ from fusion_bench import (
 from fusion_bench.dataset import CLIPDataset
 from fusion_bench.modelpool import ResNetForImageClassificationPool
 from fusion_bench.tasks.clip_classification import get_num_classes
+from torch.utils.data import random_split
 
 log = get_rankzero_logger(__name__)
 
 
 @auto_register_config
 class ImageClassificationFineTuning(BaseAlgorithm):
+    """Fine-tuning algorithm for image classification models.
+
+    This class implements end-to-end fine-tuning for image classification tasks using PyTorch Lightning.
+    It supports both epoch-based and step-based training with configurable optimizers, learning rate
+    schedulers, and data loaders.
+
+    Args:
+        max_epochs (Optional[int]): Maximum number of training epochs. Mutually exclusive with max_steps.
+        max_steps (Optional[int]): Maximum number of training steps. Mutually exclusive with max_epochs.
+        label_smoothing (float): Label smoothing factor for cross-entropy loss (0.0 = no smoothing).
+        optimizer (DictConfig): Configuration for the optimizer (e.g., Adam, SGD).
+        lr_scheduler (DictConfig): Configuration for the learning rate scheduler.
+        dataloader_kwargs (DictConfig): Additional keyword arguments for DataLoader construction.
+        **kwargs: Additional arguments passed to the base class.
+
+    Raises:
+        AssertionError: If both max_epochs and max_steps are provided.
+
+    Example:
+        ```python
+        >>> config = {
+        ...     'max_epochs': 10,
+        ...     'max_steps': None,
+        ...     'label_smoothing': 0.1,
+        ...     'optimizer': {'_target_': 'torch.optim.Adam', 'lr': 0.001},
+        ...     'lr_scheduler': {'_target_': 'torch.optim.lr_scheduler.StepLR', 'step_size': 5},
+        ...     'dataloader_kwargs': {'batch_size': 32, 'num_workers': 4}
+        ... }
+        >>> algorithm = ImageClassificationFineTuning(**config)
+        ```
+    """
+
     def __init__(
         self,
         max_epochs: Optional[int],
         max_steps: Optional[int],
+        training_data_ratio: Optional[float],
         label_smoothing: float,
         optimizer: DictConfig,
         lr_scheduler: DictConfig,
         dataloader_kwargs: DictConfig,
+        save_top_k: int,
+        save_interval: int,
+        save_on_train_epoch_end: bool,
         **kwargs,
     ):
         super().__init__(**kwargs)
-        assert (max_epochs is None) or (
+        assert (max_epochs is None or max_epochs < 0) or (
             max_steps is None or max_steps < 0
         ), "Only one of max_epochs or max_steps should be set."
-        self.training_interval = "epoch" if max_epochs is not None else "step"
+        self.training_interval = (
+            "epoch" if max_epochs is not None and max_epochs > 0 else "step"
+        )
         if self.training_interval == "epoch":
             self.max_steps = -1
         log.info(f"Training interval: {self.training_interval}")
         log.info(f"Max epochs: {max_epochs}, max steps: {max_steps}")
 
     def run(self, modelpool: ResNetForImageClassificationPool):
+        """Execute the fine-tuning process on the provided model pool.
+
+        This method performs the complete fine-tuning workflow:
+        1. Loads the pretrained model from the model pool
+        2. Prepares training and validation datasets
+        3. Configures optimizer and learning rate scheduler
+        4. Sets up Lightning trainer with appropriate callbacks
+        5. Executes the training process
+        6. Saves the final fine-tuned model
+        """
         # load model and dataset
         model = modelpool.load_pretrained_or_first_model()
         assert isinstance(model, nn.Module), "Loaded model is not a nn.Module."
@@ -59,7 +114,17 @@ class ImageClassificationFineTuning(BaseAlgorithm):
         ), "Exactly one training dataset is required."
         self.dataset_name = dataset_name = modelpool.train_dataset_names[0]
         num_classes = get_num_classes(dataset_name)
+        log.info(f"Number of classes for dataset {dataset_name}: {num_classes}")
         train_dataset = modelpool.load_train_dataset(dataset_name)
+        log.info(f"Training dataset size: {len(train_dataset)}")
+        if self.training_data_ratio is not None and 0 < self.training_data_ratio < 1:
+            train_dataset, _ = random_split(
+                train_dataset,
+                lengths=[self.training_data_ratio, 1 - self.training_data_ratio],
+            )
+            log.info(
+                f"Using {len(train_dataset)} samples for training after applying training_data_ratio={self.training_data_ratio}."
+            )
         train_dataset = CLIPDataset(
             train_dataset, processor=modelpool.load_processor(stage="train")
         )
@@ -91,7 +156,11 @@ class ImageClassificationFineTuning(BaseAlgorithm):
             objective=nn.CrossEntropyLoss(label_smoothing=self.label_smoothing),
             metrics={
                 "acc@1": Accuracy(task="multiclass", num_classes=num_classes),
-                "acc@5": Accuracy(task="multiclass", num_classes=num_classes, top_k=5),
+                f"acc@{min(5,num_classes)}": Accuracy(
+                    task="multiclass",
+                    num_classes=num_classes,
+                    top_k=min(5, num_classes),
+                ),
             },
         )
 
@@ -108,11 +177,21 @@ class ImageClassificationFineTuning(BaseAlgorithm):
             callbacks=[
                 pl_callbacks.LearningRateMonitor(logging_interval="step"),
                 pl_callbacks.DeviceStatsMonitor(),
+                pl_callbacks.ModelCheckpoint(
+                    save_top_k=self.save_top_k,
+                    every_n_train_steps=(
+                        self.save_interval if self.training_interval == "step" else None
+                    ),
+                    every_n_epochs=(
+                        self.save_interval
+                        if self.training_interval == "epoch"
+                        else None
+                    ),
+                    save_on_train_epoch_end=self.save_on_train_epoch_end,
+                    save_last=True,
+                ),
             ],
-            logger=TensorBoardLogger(
-                save_dir=log_dir,
-                name="",
-            ),
+            logger=TensorBoardLogger(save_dir=log_dir, name="", version=""),
             fast_dev_run=RuntimeConstants.debug,
         )
 
@@ -129,10 +208,26 @@ class ImageClassificationFineTuning(BaseAlgorithm):
                     "raw_checkpoints",
                     "final",
                 ),
+                algorithm_config=self.config,
+                description=f"Fine-tuned ResNet model on dataset {dataset_name}.",
             )
         return model
 
     def get_dataloader(self, dataset, stage: str):
+        """Create a DataLoader for the specified dataset and training stage.
+
+        Constructs a PyTorch DataLoader with stage-appropriate configurations:
+        - Training stage: shuffling enabled by default
+        - Validation/test stages: shuffling disabled by default
+
+        Args:
+            dataset: The dataset to wrap in a DataLoader.
+            stage (str): Training stage, must be one of "train", "val", or "test".
+                Determines default shuffling behavior.
+
+        Returns:
+            DataLoader: Configured DataLoader for the given dataset and stage.
+        """
         assert stage in ["train", "val", "test"], f"Invalid stage: {stage}"
         dataloader_kwargs = dict(self.dataloader_kwargs)
         if "shuffle" not in dataloader_kwargs:
@@ -142,10 +237,42 @@ class ImageClassificationFineTuning(BaseAlgorithm):
 
 @auto_register_config
 class ImageClassificationFineTuning_Test(BaseAlgorithm):
+    """Test/evaluation algorithm for fine-tuned image classification models.
+
+    This class implements model evaluation on test or validation datasets using PyTorch Lightning.
+    It can either evaluate a model directly or load a model from a checkpoint before evaluation.
+    The evaluation computes standard classification metrics including top-1 and top-5 accuracy.
+
+    Args:
+        checkpoint_path (str): Path to the model checkpoint file. If None, uses the model
+            directly from the model pool without loading from checkpoint.
+        dataloader_kwargs (DictConfig): Additional keyword arguments for DataLoader construction.
+        **kwargs: Additional arguments passed to the base class.
+
+    Example:
+        ```python
+        >>> config = {
+        ...     'checkpoint_path': '/path/to/model/checkpoint.ckpt',
+        ...     'dataloader_kwargs': {'batch_size': 64, 'num_workers': 4}
+        ... }
+        >>> test_algorithm = ImageClassificationFineTuning_Test(**config)
+        ```
+    """
+
     def __init__(self, checkpoint_path: str, dataloader_kwargs: DictConfig, **kwargs):
         super().__init__(**kwargs)
 
-    def run(self, modelpool: BaseModelPool):
+    def run(self, modelpool: ResNetForImageClassificationPool):
+        """Execute model evaluation on the provided model pool's test/validation dataset.
+
+        This method performs the complete evaluation workflow:
+        1. Loads the model from the model pool (pretrained or first available)
+        2. Prepares the test or validation dataset (prioritizes test if both available)
+        3. Sets up the Lightning module with appropriate metrics (top-1 and top-5 accuracy)
+        4. Loads from checkpoint if specified, otherwise uses the model directly
+        5. Executes the evaluation using Lightning trainer
+        6. Logs and returns the test metrics
+        """
         assert (
             modelpool.has_val_dataset or modelpool.has_test_dataset
         ), "No validation or test dataset found in the model pool."
@@ -181,8 +308,10 @@ class ImageClassificationFineTuning_Test(BaseAlgorithm):
                 model,
                 metrics={
                     "acc@1": Accuracy(task="multiclass", num_classes=num_classes),
-                    "acc@5": Accuracy(
-                        task="multiclass", num_classes=num_classes, top_k=5
+                    f"acc@{min(5,num_classes)}": Accuracy(
+                        task="multiclass",
+                        num_classes=num_classes,
+                        top_k=min(5, num_classes),
                     ),
                 },
             )
@@ -192,8 +321,10 @@ class ImageClassificationFineTuning_Test(BaseAlgorithm):
                 model=model,
                 metrics={
                     "acc@1": Accuracy(task="multiclass", num_classes=num_classes),
-                    "acc@5": Accuracy(
-                        task="multiclass", num_classes=num_classes, top_k=5
+                    f"acc@{min(5,num_classes)}": Accuracy(
+                        task="multiclass",
+                        num_classes=num_classes,
+                        top_k=min(5, num_classes),
                     ),
                 },
             )
@@ -207,6 +338,19 @@ class ImageClassificationFineTuning_Test(BaseAlgorithm):
         return model
 
     def get_dataloader(self, dataset, stage: str):
+        """Create a DataLoader for the specified dataset and evaluation stage.
+
+        Constructs a PyTorch DataLoader with stage-appropriate configurations for evaluation.
+        Similar to the training version but typically used for test/validation datasets.
+
+        Args:
+            dataset: The dataset to wrap in a DataLoader.
+            stage (str): Evaluation stage, must be one of "train", "val", or "test".
+                Determines default shuffling behavior (disabled for non-train stages).
+
+        Returns:
+            DataLoader: Configured DataLoader for the given dataset and stage.
+        """
         assert stage in ["train", "val", "test"], f"Invalid stage: {stage}"
         dataloader_kwargs = dict(self.dataloader_kwargs)
         if "shuffle" not in dataloader_kwargs:

fusion_bench/method/dare/simple_average.py

@@ -1,6 +1,6 @@
 import logging
 
-from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench import BaseAlgorithm, BaseModelPool, auto_register_config
 from fusion_bench.utils.state_dict_arithmetic import state_dict_add, state_dict_mul
 
 from .task_arithmetic import DareTaskArithmetic
@@ -8,6 +8,7 @@ from .task_arithmetic import DareTaskArithmetic
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class DareSimpleAverage(BaseAlgorithm):
 
     def __init__(
@@ -17,10 +18,10 @@ class DareSimpleAverage(BaseAlgorithm):
         rescale: bool = True,
         **kwargs,
     ):
+        super().__init__(**kwargs)
         self.sparsity_ratio = sparsity_ratio
         self.only_on_linear_weight = only_on_linear_weights
         self.rescale = rescale
-        super().__init__(**kwargs)
 
     def run(self, modelpool: BaseModelPool):
         return DareTaskArithmetic(

fusion_bench/method/dare/task_arithmetic.py

@@ -1,7 +1,7 @@
 import torch
 from torch import Tensor, nn
 
-from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench import BaseAlgorithm, BaseModelPool, auto_register_config
 from fusion_bench.utils.state_dict_arithmetic import state_dict_sum
 
 from .utils import (
@@ -12,6 +12,7 @@ from .utils import (
 )
 
 
+@auto_register_config
 class DareTaskArithmetic(BaseAlgorithm):
     """
     Implementation of Task Arithmetic w/ DARE.
@@ -27,11 +28,11 @@ class DareTaskArithmetic(BaseAlgorithm):
         rescale: bool = True,
         **kwargs,
     ):
+        super().__init__(**kwargs)
         self.scaling_factor = scaling_factor
         self.sparsity_ratio = sparsity_ratio
         self.only_on_linear_weights = only_on_linear_weights
         self.rescale = rescale
-        super().__init__(**kwargs)
 
     def _load_task_vector(
         self,

fusion_bench/method/dop/__init__.py

@@ -0,0 +1 @@
+from .dop import ContinualDOPForCLIP

fusion_bench/method/dop/dop.py

@@ -0,0 +1,366 @@
+"""
+Continual Model Merging without Data: Dual Projections for Balancing Stability and Plasticity. NeurIPS, 2025.
+
+
+Example:
+
+fusion_bench \
+    method=dop/dop \
+    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8_model_only \
+    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
+"""
+
+import logging
+import os
+import random
+from copy import deepcopy
+from pathlib import Path
+from typing import Dict, List, Literal, Optional, Tuple, cast
+
+import lightning as L
+import numpy as np
+import torch
+from omegaconf import DictConfig
+from torch import Tensor, nn
+from torch.autograd import Variable
+from tqdm.auto import tqdm
+from transformers import CLIPVisionModel
+
+from fusion_bench import BaseAlgorithm, BaseModelPool, auto_register_config
+from fusion_bench.method.simple_average import simple_average
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.taskpool import CLIPVisionModelTaskPool
+from fusion_bench.utils import seed_everything_by_time
+from fusion_bench.utils.json import save_to_json
+
+from .min_norm_solvers import MinNormSolver, gradient_normalizers
+from .utils import is_leaf_module, svd
+
+log = logging.getLogger(__name__)
+
+
+@auto_register_config
+class ContinualDOPForCLIP(BaseAlgorithm, LightningFabricMixin):
+
+    def __init__(
+        self,
+        seed: Optional[int] = None,
+        shuffle_order: bool = False,
+        save_on_every_step: bool = True,
+        evaluate_on_every_step: bool = False,
+        lr: float = 1e-4,
+        num_steps: int = 200,
+        mgda: bool = True,
+        ema: bool = True,
+        ema_beta: float = 0.99,
+        alpha: float = None,
+        svd_epsilon: float = 1.0,
+        svd_proj_space: str = "uv",
+        **kwargs,
+    ):
+        self.lr = lr
+        self.num_steps = num_steps
+        self.mgda = mgda
+        self.ema = ema
+        self.ema_beta = ema_beta
+        self.alpha = alpha
+        self.svd_epsilon = svd_epsilon
+        self.svd_proj_space = svd_proj_space
+        self.seed = seed
+        self.shuffle_order = shuffle_order
+        self.save_on_every_step = save_on_every_step
+        self.evaluate_on_every_step = evaluate_on_every_step
+
+        assert (
+            self.svd_epsilon >= 0 and self.svd_epsilon <= 1
+        ), "The svd_epsilon should be in the range of [0, 1]"
+        assert (
+            self.alpha >= 0 and self.alpha <= 1
+        ), "The alpha should be in the range of [0, 1]"
+        super().__init__(**kwargs)
+
+    def print_params(self, pretrained_model):
+        total_params = 0
+        linear_params = 0
+        linear_weight_params = 0
+        for module_name, module in pretrained_model.named_modules():
+            if not is_leaf_module(module):
+                continue
+            if isinstance(module, nn.Linear):
+                linear_params += sum(p.numel() for n, p in module.named_parameters())
+                linear_weight_params += sum(
+                    p.numel() for n, p in module.named_parameters() if "weight" in n
+                )
+            total_params += sum(p.numel() for p in module.parameters())
+
+        linear_ratio = linear_params / total_params * 100
+        linear_weight_ratio = linear_weight_params / total_params * 100
+        print(f"Total Parameters: {total_params}")
+        print(f"Linear Parameters: {linear_params}")
+        print(f"Linear Weight Parameters: {linear_weight_params}")
+        print(f"Linear Ratio: {linear_ratio:.2f}%")
+        print(f"Linear Weight Ratio: {linear_weight_ratio:.2f}%")
+
+    def run(self, modelpool: BaseModelPool):
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+        else:
+            seed_everything_by_time(self.fabric)
+
+        # get the model names, shuffle if needed
+        # the model names will be saved to the log directory as `model_names.json`
+        model_names = modelpool.model_names
+        if self.shuffle_order:
+            random.shuffle(model_names)
+        if self.log_dir is not None:
+            save_to_json(model_names, os.path.join(self.log_dir, "model_names.json"))
+
+        if self.evaluate_on_every_step:
+            """Configuration for the test datasets"""
+            self.taskpool = cast(CLIPVisionModelTaskPool, self._program.taskpool)
+            self._test_datasets = deepcopy(self.taskpool._test_datasets)
+
+        pretrained_model = modelpool.load_pretrained_model()
+
+        merged_model = None
+        for model_idx, model_name in enumerate(model_names):
+            print(
+                f"--------- Optimizing {model_idx + 1}/{len(model_names)}-th with {model_name} ---------"
+            )
+            if model_idx == 0:
+                merged_model = modelpool.load_model(model_names[0])
+            else:
+                merged_model = self._layer_wise_optimize(
+                    model_names=["merged", model_name],
+                    pretrained_model=deepcopy(pretrained_model),
+                    finetuned_models={
+                        "merged": merged_model,
+                        model_name: modelpool.load_model(model_name),
+                    },
+                    model_idx=model_idx,
+                )
+
+            if self.save_on_every_step:
+                self.save_merged_model(merged_model, model_idx)
+
+            if self.evaluate_on_every_step:
+                self.taskpool._is_setup = False
+                self.taskpool._test_datasets = DictConfig(
+                    {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
+                )
+                report = self.taskpool.evaluate(deepcopy(merged_model))
+                save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
+
+        return merged_model
+
+    def _layer_wise_optimize(
+        self,
+        model_names: List[str],
+        pretrained_model: nn.Module,
+        finetuned_models: Dict[str, nn.Module],
+        model_idx: int,
+    ):
+        time_cost = []
+        for module_name, module in pretrained_model.named_modules():
+            if not is_leaf_module(module):
+                continue
+
+            if isinstance(module, nn.Linear):
+                if module.weight.requires_grad:
+                    import time
+
+                    start_time = time.time()
+                    merged_weight = self._optimize_weight(
+                        module.weight,
+                        {
+                            model_name: finetuned_models[model_name]
+                            .get_submodule(module_name)
+                            .weight
+                            for model_name in model_names
+                        },
+                        module_name,
+                        model_idx,
+                    )
+                    end_time = time.time()
+                    time_cost.append(end_time - start_time)
+                    module.weight.data = merged_weight.data
+                else:
+                    module.weight.data = simple_average(
+                        [
+                            finetuned_models[model_name]
+                            .get_submodule(module_name)
+                            .weight
+                            for model_name in model_names
+                        ]
+                    )
+                if module.bias is not None:
+                    module.bias.data = simple_average(
+                        [
+                            finetuned_models[model_name].get_submodule(module_name).bias
+                            for model_name in model_names
+                        ]
+                    )
+            else:
+                simple_average(
+                    [
+                        finetuned_models[model_name].get_submodule(module_name)
+                        for model_name in model_names
+                    ],
+                    base_module=module,
+                )
+
+        return pretrained_model
+
+    def _optimize_weight(
+        self,
+        pretrained_weight: Tensor,
+        finetuned_weights: Dict[str, Tensor],
+        module_name: str,
+        model_idx: int,
+    ):
+        assert (
+            self.fabric.world_size == 1
+        ), "This algorithm is not currently supported in distributed training"
+
+        pretrained_weight = self.fabric.to_device(pretrained_weight.detach())
+        finetuned_weights = {
+            model_name: self.fabric.to_device(finetuned_weight.detach())
+            for model_name, finetuned_weight in finetuned_weights.items()
+        }
+
+        merged_weight = self.fabric.to_device(
+            nn.Parameter(
+                simple_average(
+                    [
+                        finetuned_weight.detach()
+                        for finetuned_weight in finetuned_weights.values()
+                    ]
+                ),
+                requires_grad=True,
+            )
+        )
+
+        # Compute SVD of the difference between the finetuned and pretrained weights
+        proj_u_dict = {}
+        proj_v_dict = {}
+        proj_s_dict = {}
+        for i, finetuned_weight in enumerate(finetuned_weights.values()):
+            finetuned_tv = finetuned_weight - pretrained_weight
+            u, s, v = svd(finetuned_tv, full_matrices=True)
+            epsilon = 1.0 if self.svd_epsilon > 1.0 else self.svd_epsilon
+            cumsum_ratio = s.cumsum(dim=0) / s.sum()
+            split_rank = torch.searchsorted(cumsum_ratio, epsilon).item()
+            u_main = u[:, :split_rank]
+            v_main = v[:, :split_rank]
+            s_main = s[:split_rank]
+            proj_u_dict[i] = u_main
+            proj_v_dict[i] = v_main
+            proj_s_dict[i] = s_main
+
+        if self.mgda:
+            if self.ema:
+                ema_sol = [self.alpha, 1 - self.alpha]
+            # This is multiple-gradient descent algorithm (MGDA) optimization
+            optimizer = torch.optim.Adam([merged_weight], lr=self.lr)
+            all_losses = [[], []]
+            all_alphas = [[], []]
+            for step_idx in tqdm(
+                range(self.num_steps), desc=f"Optimizing {module_name} weight"
+            ):
+                # Scaling the loss functions based on the algorithm choice
+                loss_data = {}
+                grads = {}
+                for i, finetuned_weight in enumerate(finetuned_weights.values()):
+                    proj_u = proj_u_dict[i]
+                    proj_v = proj_v_dict[i]
+                    proj_s = proj_s_dict[i]
+                    delta_tv = merged_weight - finetuned_weight
+                    loss_i = self.cal_loss_i(delta_tv, proj_s, proj_u, proj_v)
+                    loss_data[i] = float(loss_i.data)
+
+                    all_losses[i].append(float(loss_i.data))
+
+                    optimizer.zero_grad()
+                    loss_i.backward()
+                    grads[i] = Variable(
+                        merged_weight.grad.data.clone(), requires_grad=False
+                    )
+
+                # Normalize all gradients
+                gn = gradient_normalizers(
+                    grads=grads, losses=loss_data, normalization_type="loss"
+                )
+                for i, _ in enumerate(finetuned_weights.values()):
+                    grads[i] = grads[i] / float(gn[i])
+
+                # Frank-Wolfe iteration to compute scales.
+                sol, min_norm = MinNormSolver.find_min_norm_element(
+                    [[grads[i]] for i in range(len(finetuned_weights.values()))]
+                )
+
+                if self.ema:
+                    ema_sol = [
+                        self.ema_beta * ema_sol[i] + (1 - self.ema_beta) * float(sol[i])
+                        for i in range(len(sol))
+                    ]
+                    sol = ema_sol
+                    all_alphas[0].append(ema_sol[0])
+                    all_alphas[1].append(ema_sol[1])
+
+                # Scaled back-propagation
+                loss = 0
+                for i, finetuned_weight in enumerate(finetuned_weights.values()):
+                    # Comptue gradients of each loss function wrt parameters
+                    proj_u = proj_u_dict[i]
+                    proj_v = proj_v_dict[i]
+                    proj_s = proj_s_dict[i]
+                    delta_tv = merged_weight - finetuned_weight
+                    loss_i = self.cal_loss_i(delta_tv, proj_s, proj_u, proj_v)
+                    loss += float(sol[i]) * loss_i
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        else:
+            # This is a naive weighted optimization
+            optimizer = torch.optim.Adam([merged_weight], lr=self.lr)
+            for step_idx in tqdm(
+                range(self.num_steps), desc=f"Optimizing {module_name} weight"
+            ):
+                loss = 0
+                for i, finetuned_weight in enumerate(finetuned_weights.values()):
+                    proj_u = proj_u_dict[i]
+                    proj_v = proj_v_dict[i]
+                    proj_s = proj_s_dict[i]
+                    delta_tv = merged_weight - finetuned_weight
+                    loss_i = self.cal_loss_i(delta_tv, proj_s, proj_u, proj_v)
+                    loss += self.alpha * loss_i if i == 0 else (1 - self.alpha) * loss_i
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        return merged_weight.detach().cpu()
+
+    def cal_loss_i(self, delta_tv, proj_s, proj_u, proj_v):
+        proj_delta_1 = torch.diag(proj_s) @ proj_u.T @ delta_tv
+        proj_delta_2 = delta_tv @ proj_v @ torch.diag(proj_s)
+        loss_i_u = torch.linalg.matrix_norm(proj_delta_1, ord="fro") ** 2
+        loss_i_v = torch.linalg.matrix_norm(proj_delta_2, ord="fro") ** 2
+        if self.svd_proj_space == "uv":
+            loss_i = loss_i_u + loss_i_v
+        elif self.svd_proj_space == "u":
+            loss_i = loss_i_u
+        elif self.svd_proj_space == "v":
+            loss_i = loss_i_v
+        else:
+            raise ValueError("Invalid svd_proj_space")
+
+        return loss_i
+
+    def save_merged_model(self, merged_model: CLIPVisionModel, step: int):
+        os.makedirs(Path(self.log_dir) / "checkpoints", exist_ok=True)
+        merged_model.save_pretrained(
+            Path(self.log_dir) / "checkpoints" / f"merged_model_{step}"
+        )