PWEMoE: Pareto-Driven Weight-Ensembling Mixture of Experts

Overview of PWE MoE (a) An illustration of Pareto front learning in MOOP. Where \(P_1\) and \(P_2\) are performance metrics for two tasks, colored lines represent different Pareto optimal solutions, and the solid black line represents the Pareto front. (b) An overview of the model up-scaling process. We upcycle the MLP modules to MoE modules and merge the remaining parts using task arithmetic. (c) The MoE module, comprising a routing network and a parameter decoder network. The routing network accepts a user preference vector and generates routing weights for weight-ensembling.

Abstract

Solving multi-objective optimization problems for large deep neural networks is a challenging task due to the complexity of the loss landscape and the expensive computational cost of training and evaluating models. Efficient Pareto front approximation of large models enables multi-objective optimization for various tasks such as multi-task learning and trade-off analysis. Existing algorithms for learning Pareto set, including (1) evolutionary, hypernetworks, and hypervolume-maximization methods, are computationally expensive and have restricted scalability to large models; (2) Scalarization algorithms, where a separate model is trained for each objective ray, which is inefficient for learning the entire Pareto set and fails to capture the objective trade-offs effectively. Inspired by the recent success of model merging, we propose a practical and scalable approach to Pareto set learning problem via mixture of experts (MoE) based model fusion. By ensembling the weights of specialized single-task models, the MoE module can effectively capture the trade-offs between multiple objectives and closely approximate the entire Pareto set of large neural networks. Once the routers are learned and a preference vector is set, the MoE module can be unloaded, thus no additional computational cost is introduced during inference. We conduct extensive experiments on vision and language tasks using large-scale models such as CLIP-ViT and GPT-2. The experimental results demonstrate that our method efficiently approximates the entire Pareto front of large models. Using only hundreds of trainable parameters of the MoE routers, our method even has lower memory usage compared to linear scalarization and algorithms that learn a single Pareto optimal solution, and are scalable to both the number of objectives and the size of the model. Our method significantly reduces the computational burden of learning the Pareto set, for example, in the two-task case, it can be achieved in just a few minutes. Code is available at: GitHub .

Examples

Not tested yet

The examples provided below have not been tested yet.

For a thoroughly tested and verified implementation of the algorithm, please refer to the original repository: tanganke/pareto_set_learning . Additionally, the experimental results and further insights into the algorithm can be found in the original research paper: arXiv:2406.09770 .

PWEMoE-LS on eight image classification tasks using CLIP-ViT-B/32 models, and the results are logged to outputs/logs/ViT-B-32/PWEMoE-LS-8tasks.

fusion_bench \
    method=pwe_moe_ls_for_clip \
    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8 \
    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8 \
    fabric.loggers.root_dir=outputs/logs/ViT-B-32 \
    fabric.loggers.name=PWEMoE-LS-8tasks

References

clip_pwe_moe

PWEMoEAlgorithmForCLIP

Bases: BaseAlgorithm, SimpleProfilerMixin, CLIPClassificationMixin

Source code in fusion_bench/method/pwe_moe/clip_pwe_moe.py

class PWEMoEAlgorithmForCLIP(
    BaseAlgorithm,
    SimpleProfilerMixin,
    CLIPClassificationMixin,
):
    modelpool: CLIPVisionModelPool = None
    _config_mapping = BaseAlgorithm._config_mapping | {
        "upscale_mlp": "upscale_mlp",
        "upscale_attn": "upscale_attn",
        "init_lambda": "init_lambda",
        "router_hidden_layers": "router_hidden_layers",
        "lr": "lr",
        "num_steps": "num_steps",
        "save_interval": "save_interval",
        "alpha": "alpha",
        "checkpoint_path": "checkpoint_path",
        "eval_grid": "eval_grid",
        "eval_grid_n": "eval_grid_n",
        "eval_grid_m": "eval_grid_m",
        "_dataloader_kwargs": "dataloader_kwargs",
    }

    def __init__(
        self,
        *,
        upscale_mlp: bool,
        upscale_attn: bool,
        init_lambda: float,
        router_hidden_layers: int,
        lr: float,
        num_steps: int,
        save_interval: int,
        alpha: float,
        checkpoint_path: str,
        eval_grid: bool,
        eval_grid_n: int,
        eval_grid_m: int,
        dataloader_kwargs: DictConfig,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.upscale_mlp = upscale_mlp
        self.upscale_attn = upscale_attn
        self.init_lambda = init_lambda
        self.router_hidden_layers = router_hidden_layers
        self.lr = lr
        self.num_steps = num_steps
        self.save_interval = save_interval
        self.alpha = alpha
        self.checkpoint_path = checkpoint_path
        self.eval_grid = eval_grid
        self.eval_grid_n = eval_grid_n
        self.eval_grid_m = eval_grid_m
        self._dataloader_kwargs = dataloader_kwargs

    @override
    def run(self, modelpool: CLIPVisionModelPool):
        self.modelpool = modelpool

        model = self.setup_model()
        if self.checkpoint_path is not None:
            model.load_state_dict(torch.load(self.checkpoint_path, map_location="cpu"))
        else:
            train_loaders = self.setup_train_loaders()
            model = self.train(model, train_loaders)

        if self.eval_grid:
            return map(
                lambda m, r: {
                    "model": ParetoWeightEnsemblingModule.set_preferenece_vector(
                        m,
                        torch.as_tensor(
                            r, device=self.fabric.device, dtype=torch.float32
                        ),
                    ),
                    "preference_vector": r,
                },
                itertools.cycle([model]),
                generate_simplex_grid(self.eval_grid_n, self.eval_grid_m),
            )
        return model

    def load_clip_models(self):
        """
        Loads the pretrained CLIP model and the fine-tuned models for each dataset specified in the configuration.
        """
        # load pretrained and fine-tuned model
        with timeit_context():
            log.info("load models")
            pretrained_model: CLIPVisionModel = self.modelpool.load_model(
                "_pretrained_"
            )
            finetuned_models = {
                model_name: self.modelpool.load_model(model_name)
                for model_name in self.modelpool.model_names
            }

        log.info("pretrained model statistics:")
        print_parameters(pretrained_model)
        return pretrained_model, finetuned_models

    def setup_model(self):
        pretrained_model, finetuned_models = self.load_clip_models()
        self.setup_zero_shot_classification_head()

        with timeit_context("Building PWEMoE model"):
            model = deepcopy(pretrained_model)

            # merge the remaining layers using task arithmetic
            if self.init_lambda != 0:
                task_arithmetic_merge(
                    model,
                    finetuned_models.values(),
                    scaling_factor=self.init_lambda,
                    inplace=True,
                )
            # fix all parameters
            model.requires_grad_(False)

            num_layers = len(model.vision_model.encoder.layers)

            def get_layer(m, i):
                return cast(CLIPEncoderLayer, m.vision_model.encoder.layers[i])

            for layer_idx in tqdm(range(num_layers)):
                if self.upscale_mlp:
                    # upscale the mlp layer
                    get_layer(model, layer_idx).mlp = ParetoWeightEnsemblingModule(
                        base_model=get_layer(pretrained_model, layer_idx).mlp,
                        expert_models=[
                            get_layer(m, layer_idx).mlp
                            for m in finetuned_models.values()
                        ],
                        init_lambda=self.init_lambda,
                        fix_base_model_and_experts=True,
                        router_hidden_layers=self.router_hidden_layers,
                    )

                if self.upscale_attn:
                    # upscale the Attention layer
                    get_layer(model, layer_idx).self_attn = (
                        ParetoWeightEnsemblingModule(
                            base_model=get_layer(pretrained_model, layer_idx).self_attn,
                            expert_models=[
                                get_layer(m, layer_idx).self_attn
                                for m in finetuned_models.values()
                            ],
                            init_lambda=self.init_lambda,
                            fix_base_model_and_experts=True,
                            router_hidden_layers=self.router_hidden_layers,
                        )
                    )

            print("model statistics after upscaling:")
            print_parameters(model)
            return model

    def setup_train_loaders(self):
        """
        Loads the datasets specified in the configuration.
        """
        train_datasets = {
            dataset_name: self.modelpool.load_train_dataset(
                dataset_name, self.clip_processor
            )
            for dataset_name in self.modelpool.model_names
        }
        train_loaders = {
            dataset_name: DataLoader(dataset, shuffle=True, **self._dataloader_kwargs)
            for dataset_name, dataset in train_datasets.items()
        }
        train_loaders = {
            dataset_name: self.fabric.setup_dataloaders(loader)
            for dataset_name, loader in train_loaders.items()
        }
        return train_loaders

    def train(self, model: nn.Module, train_loaders: Dict[str, DataLoader]):
        config = self.config

        # save the configuration
        self.log_hyperparams(config, filename="method_config.yaml")

        # setup the model
        num_objectives = len(self.modelpool.model_names)
        model = model

        # setup data loaders
        train_loaders = {
            name: InfiniteDataLoader(loader) for name, loader in train_loaders.items()
        }

        # set up the optimizer and learning rate scheduler
        optimizer = torch.optim.Adam(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=config.lr,
        )
        model, optimizer = self.fabric.setup(model, optimizer)
        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer=optimizer, T_max=config.num_steps, eta_min=config.lr * 0.1
        )

        model.train()
        device = self.fabric.device
        for step_idx in tqdm(
            range(1, 1 + config.num_steps), "training", dynamic_ncols=True
        ):
            # sample a preference ray
            ray = torch.from_numpy(
                np.random.dirichlet((config.alpha,) * num_objectives, 1)
                .astype(np.float32)
                .flatten()
            ).to(device)
            ParetoWeightEnsemblingModule.set_preferenece_vector(model, ray)

            losses = []
            for dataset_idx, dataset_name in enumerate(train_loaders):
                batch = next(train_loaders[dataset_name])
                images, labels = batch

                logits = self.compute_logits(model, images, dataset_name)
                _loss = F.cross_entropy(logits, labels)
                losses.append(_loss)

            loss = self.compute_loss(model, ray, losses)

            optimizer.zero_grad()
            self.fabric.backward(loss)
            optimizer.step()

            lr_scheduler.step()

            self.fabric.log("train/loss", loss.item(), step=step_idx)

            if step_idx % config.save_interval == 0:
                (Path(self.log_dir) / "checkpoints").mkdir(exist_ok=True, parents=True)
                save_path = (
                    Path(self.log_dir) / "checkpoints" / f"model_step={step_idx}.pt"
                )
                torch.save(model.state_dict(), save_path)

        return model

    @abstractmethod
    def compute_loss(
        self, model: nn.Module, ray: Tensor, losses: List[Tensor]
    ) -> Tensor:
        """
        Computes the overall losses using the given preference ray.

        Args:
            model (nn.Module): The model being trained.
            ray (Tensor): A tensor representing the preference ray, which contains the weights for each objective.
            losses (List[Tensor]): A list of loss values for each objective.
        """
        pass

compute_loss(model, ray, losses) abstractmethod

Computes the overall losses using the given preference ray.

Parameters:

• model (Module) –

  The model being trained.

• ray (Tensor) –

  A tensor representing the preference ray, which contains the weights for each objective.

• losses (List[Tensor]) –

  A list of loss values for each objective.

Source code in fusion_bench/method/pwe_moe/clip_pwe_moe.py

@abstractmethod
def compute_loss(
    self, model: nn.Module, ray: Tensor, losses: List[Tensor]
) -> Tensor:
    """
    Computes the overall losses using the given preference ray.

    Args:
        model (nn.Module): The model being trained.
        ray (Tensor): A tensor representing the preference ray, which contains the weights for each objective.
        losses (List[Tensor]): A list of loss values for each objective.
    """
    pass

load_clip_models()

Loads the pretrained CLIP model and the fine-tuned models for each dataset specified in the configuration.

Source code in fusion_bench/method/pwe_moe/clip_pwe_moe.py

def load_clip_models(self):
    """
    Loads the pretrained CLIP model and the fine-tuned models for each dataset specified in the configuration.
    """
    # load pretrained and fine-tuned model
    with timeit_context():
        log.info("load models")
        pretrained_model: CLIPVisionModel = self.modelpool.load_model(
            "_pretrained_"
        )
        finetuned_models = {
            model_name: self.modelpool.load_model(model_name)
            for model_name in self.modelpool.model_names
        }

    log.info("pretrained model statistics:")
    print_parameters(pretrained_model)
    return pretrained_model, finetuned_models

setup_train_loaders()

Loads the datasets specified in the configuration.

Source code in fusion_bench/method/pwe_moe/clip_pwe_moe.py

def setup_train_loaders(self):
    """
    Loads the datasets specified in the configuration.
    """
    train_datasets = {
        dataset_name: self.modelpool.load_train_dataset(
            dataset_name, self.clip_processor
        )
        for dataset_name in self.modelpool.model_names
    }
    train_loaders = {
        dataset_name: DataLoader(dataset, shuffle=True, **self._dataloader_kwargs)
        for dataset_name, dataset in train_datasets.items()
    }
    train_loaders = {
        dataset_name: self.fabric.setup_dataloaders(loader)
        for dataset_name, loader in train_loaders.items()
    }
    return train_loaders

PWEMoELinearScalarizationForCLIP

Bases: PWEMoEAlgorithmForCLIP

Source code in fusion_bench/method/pwe_moe/clip_pwe_moe.py

class PWEMoELinearScalarizationForCLIP(PWEMoEAlgorithmForCLIP):
    def compute_loss(self, model, ray, losses):
        loss = 0
        for r, l in zip(ray, losses):
            loss += r * l
        return loss

PWEMoExactParetoOptimalForCLIP

Bases: PWEMoEAlgorithmForCLIP

Source code in fusion_bench/method/pwe_moe/clip_pwe_moe.py

class PWEMoExactParetoOptimalForCLIP(PWEMoEAlgorithmForCLIP):
    def compute_loss(self, model: nn.Module, ray: Tensor, losses: Tuple[Tensor]):
        from phn.solvers import EPOSolver

        if self.epo_solver is None:
            num_objectives = len(self.finetuned_models)
            self.epo_solver = EPOSolver(n_tasks=num_objectives, n_params=None)
        epo_solver = self.epo_solver

        losses = torch.stack(losses)
        loss = epo_solver.get_weighted_loss(
            losses,
            ray,
            tuple(filter(lambda p: p.requires_grad, model.parameters())),
        )
        return loss