+ In the complex landscape of multi-task learning, AdaMerging has emerged as a potent method for adaptively merging model parameters to optimize performance across tasks. Unlike traditional fixed-coefficient methods, AdaMerging autonomously learns merging coefficients, offering a more refined and responsive approach1.
+The cornerstone of AdaMerging lies in its adaptive nature, where it learns the coefficients for merging either on a task-wise or layer-wise basis. This adaptability is driven by an entropy minimization strategy applied to unlabeled test samples as a surrogate objective function, which serves to refine the merging coefficients for optimal performance.
+Task-wise AdaMerging is formulated as:
+where \(\lambda_i\) represents the merging coefficient for the \(i\)-th task, and \(\tau_i\) denotes the task vector for the \(i\)-th task.
+On the other hand, Layer-wise AdaMerging is articulated as:
+where the merging coefficient \(\lambda^{l}_{i}\) and task vector \(\tau^{l}_{i}\) are specific to each layer \(l\) of the model.
+By leveraging this adaptive learning approach, AdaMerging significantly enhances the model's ability to generalize across tasks and layers, resulting in a more robust and finely-tuned performance profile. The method’s reliance on entropy minimization ensures that the merging process continually seeks the most informative and stable configuration, adapting to the specific needs of the dataset and tasks at hand.
+Task-wise Coefficients. +The below Figure shows the changes during the iteration process of merging coefficient optimization of each task vector in Task-wise AdaMerging and AdaMerging++, which is shown every ten steps. We consistently observe that the merging coefficients of each task vector are inconsistent. When the number of tasks is relatively large, it is obviously undesirable to grid search the coefficients of each task, but our AdaMerging avoids this manual search process.
+
+Layer-wise Coefficients. +The following Figure shows the merging coefficients learned by Layer-wise AdaMerging and AdaMerging++ on ViT-B/32 respectively. We observed that:
+
+Merge CLIP-ViT-B/32 models from eight downstream image classification tasks:
+fusion_bench \
+ method=adamerging \
+ method.name=clip_layer_wise_adamerging \
+ method.save_merging_weights=merging_weights.pt \
+ modelpool=clip-vit-base-patch32_TA8 \
+ taskpool=clip-vit-classification_TA8 \
+ fabric.loggers.root_dir=outputs/logs/ViT-B-32 \
+ fabric.loggers.name=clip_layer_wise_adamerging_adam
+Part of the output:
+Profiler Report
+
+----------------------------------------------------------------------------------------------------------------------------------
+| Action | Mean duration (s) | Num calls | Total time (s) | Percentage % |
+----------------------------------------------------------------------------------------------------------------------------------
+| Total | - | 26001 | 724.65 | 100 % |
+----------------------------------------------------------------------------------------------------------------------------------
+| backward pass | 0.060172 | 8000 | 481.38 | 66.429 |
+| forward pass | 0.016124 | 8000 | 128.99 | 17.801 |
+| data loading | 0.0063443 | 8000 | 50.754 | 7.004 |
+| merging weights | 0.050735 | 1000 | 50.735 | 7.0013 |
+| construct the wrapped model | 7.2558 | 1 | 7.2558 | 1.0013 |
+| optimizer step | 0.00098186 | 1000 | 0.98186 | 0.13549 |
+----------------------------------------------------------------------------------------------------------------------------------
+task_wise_adamerging
+
+
+¶TaskWiseAdaMergingAlgorithm
+
+
+¶
+ Bases: ModelFusionAlgorithm
fusion_bench/method/adamerging/task_wise_adamerging.py38 + 39 + 40 + 41 + 42 + 43 + 44 + 45 + 46 + 47 + 48 + 49 + 50 + 51 + 52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 | |
compute_logits(module, batch, task)
+
+
+ abstractmethod
+
+
+¶Compute the logits for the given batch and task.
+ + +Parameters:
+module
+¶Module)
+ –
+ The model module.
+batch
+¶tuple)
+ –
+ A batch of input data.
+task
+¶str)
+ –
+ The name of the task.
+Returns:
+Tensor ( Tensor
+) –
+ The classification logits for the batch.
+fusion_bench/method/adamerging/task_wise_adamerging.pyentropy_loss(logits)
+
+¶Compute the entropy loss of a set of logits.
+ + +Parameters:
+logits
+¶Tensor)
+ –
+ The logits to compute the entropy loss of.
+Returns:
+Tensor ( Tensor
+) –
+ The entropy loss of the logits.
+fusion_bench/method/adamerging/task_wise_adamerging.pyclip_task_wise_adamerging
+
+
+¶CLIPTaskWiseAdaMergingAlgorithm
+
+
+¶
+ Bases: TaskWiseAdaMergingAlgorithm
A class for task-wise adaptive merging of CLIP models.
+This class extends the TaskWiseAdaMergingAlgorithm to provide specific +functionality for CLIP models, including loading datasets, constructing +zero-shot classification heads, and computing logits.
+ + +Attributes:
+modelpool
+ (CLIPVisionModelPool)
+ –
+ The model pool containing CLIP models.
+_clip_processor
+ (CLIPProcessor)
+ –
+ The CLIP processor for preparing inputs.
+zeroshot_weights
+ (dict)
+ –
+ A dictionary to store zero-shot weights for each task.
+fusion_bench/method/adamerging/clip_task_wise_adamerging.py51 + 52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 | |
compute_logits(module, batch, task)
+
+¶Compute the logits for the given batch and task.
+This method computes the image embeddings, normalizes them, and calculates +the cosine similarity with the text embeddings to produce classification logits.
+ + +Parameters:
+module
+¶Module)
+ –
+ The model module.
+batch
+¶tuple)
+ –
+ A batch of input data.
+task
+¶str)
+ –
+ The name of the task.
+Returns:
+Tensor ( Tensor
+) –
+ The classification logits for the batch.
+fusion_bench/method/adamerging/clip_task_wise_adamerging.pyget_shuffled_test_loader_iter(task)
+
+
+ cached
+
+
+¶Get an iterator over the shuffled test DataLoader for the task.
+ + +Parameters:
+task
+¶str)
+ –
+ The name of the task.
+Returns:
+iterator –
+ An iterator over the shuffled test DataLoader.
+fusion_bench/method/adamerging/clip_task_wise_adamerging.pyget_test_dataset(task)
+
+
+ cached
+
+
+¶Load the test dataset for the task. +This method is cached, so the dataset is loaded only once.
+ + +Parameters:
+task
+¶str)
+ –
+ The name of the task.
+Returns:
+CLIPDataset –
+ The test dataset for the task.
+fusion_bench/method/adamerging/clip_task_wise_adamerging.pyon_test_time_adaptation_start()
+
+¶Prepare for test-time adaptation.
+This method loads the CLIP processor and constructs the zero-shot +classification head for each task.
+ +fusion_bench/method/adamerging/clip_task_wise_adamerging.pyInfiniteDataLoader
+
+
+¶A wrapper class for DataLoader to create an infinite data loader. +This is useful in case we are only interested in the number of steps and not the number of epochs.
+This class wraps a DataLoader and provides an iterator that resets +when the end of the dataset is reached, creating an infinite loop.
+ + +Attributes:
+data_loader
+ (DataLoader)
+ –
+ The DataLoader to wrap.
+data_iter
+ (iterator)
+ –
+ An iterator over the DataLoader.
+fusion_bench/method/adamerging/clip_task_wise_adamerging.pylayer_wise_adamerging
+
+
+¶LayerWiseAdaMergingAlgorithm
+
+
+¶
+ Bases: ModelFusionAlgorithm, LightningFabricMixin, SimpleProfilerMixin
Implements the Layer-Wise AdaMerging Algorithm.
+This class merges the layers of a pretrained model with those of several fine-tuned models. +The merging is controlled by layer-wise weights, which can be initialized based on a provided configuration or loaded from a file.
+ + + + + + +fusion_bench/method/adamerging/layer_wise_adamerging.py28 + 29 + 30 + 31 + 32 + 33 + 34 + 35 + 36 + 37 + 38 + 39 + 40 + 41 + 42 + 43 + 44 + 45 + 46 + 47 + 48 + 49 + 50 + 51 + 52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 +217 +218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 | |
__init__(algorithm_config)
+
+¶Initialize the LayerWiseAdaMergingAlgorithm with the given configuration.
+ + +Parameters:
+algorithm_config
+¶DictConfig)
+ –
+ The configuration for the algorithm.
+fusion_bench/method/adamerging/layer_wise_adamerging.pycompute_logits(module, images, task)
+
+
+ abstractmethod
+
+
+¶Compute the logits for the given images and task.
+ + +Parameters:
+module
+¶The model module.
+images
+¶Tensor)
+ –
+ The input images.
+task
+¶str)
+ –
+ The name of the task.
+Returns:
+Tensor ( Tensor
+) –
+ The computed logits.
+fusion_bench/method/adamerging/layer_wise_adamerging.pyconstruct_layer_wise_merged_model(modelpool)
+
+¶Constructs a wrapped layer-wise merged model from model pool.
+This method creates a new wrapped model by merging the layers of a pretrained model with those of several fine-tuned models.
+The merging is controlled by layer-wise weights, which is a torch.Tensor of the shape (num_models, num_layers).
+The merging weights can be initialized based on a provided configuration or loaded from a file.
Parameters:
+modelpool
+¶ModelPool)
+ –
+ An object containing the pretrained model and fine-tuned models to be merged.
+Returns:
+LayerWiseMergedModel –
+ An instance of the merged model with layer-wise weights applied.
+fusion_bench/method/adamerging/layer_wise_adamerging.pyget_shuffled_test_loader_iter(task)
+
+
+ abstractmethod
+
+
+¶Loader of test dataset for test-time adaptation. labels are not needed.
+ + +Parameters:
+task
+¶str)
+ –
+ The name of the task.
+Returns:
+DataLoader ( DataLoader
+) –
+ The data loader for the test dataset.
+fusion_bench/method/adamerging/layer_wise_adamerging.pyon_test_time_adaptation_start()
+
+¶Something to do before the test-time adaptation starts. Such as setting up the task-specific heads.
+ + +run(modelpool)
+
+¶Run the Layer-Wise AdaMerging Algorithm.
+This method constructs the wrapped model and performs test-time adaptation if necessary.
+ + +Parameters:
+modelpool
+¶ModelPool)
+ –
+ The model pool containing the pretrained and fine-tuned models.
+Returns:
+LayerWiseMergedModel –
+ The merged model after test-time adaptation.
+fusion_bench/method/adamerging/layer_wise_adamerging.pysave_merging_weights(file_path, merging_weights)
+
+¶Save the merging weights to a file.
+ + +Parameters:
+file_path
+¶str)
+ –
+ The path to save the merging weights.
+merging_weights
+¶Tensor)
+ –
+ The merging weights to save.
+fusion_bench/method/adamerging/layer_wise_adamerging.pytest_time_adaptation(module)
+
+¶Perform test-time adaptation on the merged model.
+This method adapts the merging weights during test-time to improve performance.
+ + +Parameters:
+module
+¶LayerWiseMergedModel)
+ –
+ The merged model.
+Returns:
+LayerWiseMergedModel –
+ The adapted merged model.
+fusion_bench/method/adamerging/layer_wise_adamerging.pyclip_layer_wise_adamerging
+
+
+¶Example Usage:
+fusion_bench method=adamerging method.name=clip_layer_wise_adamerging method.save_merging_weights=merging_weights.pt modelpool=clip-vit-base-patch32_TA8 taskpool=clip-vit-classification_TA8 fabric.loggers.root_dir=outputs/logs/ViT-B-32 fabric.loggers.name=clip_layer_wise_adamerging_adam
+CLIPLayerWiseAdaMergingAlgorithm
+
+
+¶
+ Bases: CLIPClassificationMixin, LayerWiseAdaMergingAlgorithm
fusion_bench/method/adamerging/clip_layer_wise_adamerging.pyon_test_time_adaptation_start()
+
+¶Here we load the CLIP processor and construct the zero-shot classification head for each task.
+ + +(ICLR 2024) AdaMerging: Adaptive Model Merging for Multi-Task Learning. https://openreview.net/pdf?id=nZP6NgD3QY ↩
+Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks? Advances in neural information processing systems, 27, 2014. ↩
+A. Tang, L. Shen, Y. Luo, N. Yin, L. Zhang, and D. Tao, “Merging Multi-Task Models via Weight-Ensembling Mixture of Experts,” ICML 2024. doi: 10.48550/arXiv.2402.00433. ↩
+
+Consider a discrete categorical distribution parameterized by logits \(\mathbf{x} = (x_1, \dots, x_n) \in \mathbb{R}^{n}\), where \(x_i\) is the logit of the \(i\)-th category. The Gumbel-Max trick 123 states a reparameterization trick to sample from the categorical distribution by sampling from the standard Gumbel distribution \(\text{Gumbel}(\mu=0,\beta=1)\) and taking the argmax of the sum of the Gumbel random variables and the logits.
+This trick proceeds as follows: +sample \(n\) Gumbel random variables \(g_1, \dots, g_n\) independently from the standard Gumbel distribution \(\text{Gumbel}(\mu=0,\beta=1)\) (We can draw a random sample \(u\) from a unifrom distribution on the interval \((0,1)\) and then transform it into a Gumbel-distributed variable \(g\) using the formula \(g=-\log(-\log u)\).), find the index \(i\) of that maximizes \(x_i + g_i\), then we have
+If we represent the categorical distribution as a one-hot vector \(\mathbf{y} = (y_1, \dots, y_n) \in \{0,1\}^n\), where \(y_i=1\) indicates that the \(i\)-th category is sampled and for all \(j\neq i\), \(y_j=0\), then we have
+Since the derivative of the \({\arg\max}\) function is not defined, we cannot backpropagate the gradients through it. +To address this issue, (Maddison et al., 2017)4 proposed to use a continuous relaxation of the discrete categorical distribution. +A CONCRETE random variable (CONtinuous relaxation of disCRETE random variable) relax the condition that the one-hot vector \(\mathbf{y}\) must be located at the vertices of the \((n-1)\)-dimensional simplex \(\Delta^{n-1}\), and instead, it allows \(\mathbf{y}\) to be located anywhere inside the simplex \(\Delta^{n-1}\), i.e. \(\{ y\in \mathbb{R}^n | y_i \in [0,1], \sum_{i=1}^n y_i =1 \}\).
+To sample a Concrete random variable \(\mathbf{y}\) from a distribution that is parameterized by a temperature hyperparameter \(\lambda > 0\) and a vector of logits \(\mathbf{x} = (x_1, \dots, x_n) \in \mathbb{R}^{n}\), we have
+where \(\mathbf{g} = (g_1, \dots, g_n)\) is a vector of Gumbel random variables that are independently sampled from the standard Gumbel distribution \(\text{Gumbel}(\mu=0,\beta=1)\).
+A subspace mask \(\mathbf{m}\) is a binary vector that identifies a subspace of the parameter space. +For a neural network parametrized by \(\theta\), we can use a subspace mask \(\mathbf{m}\) to identify a subspace of the parameter space \(\mathbf{\theta}\) by setting the parameters that are not in the subspace to zero, i.e. \(\mathbf{\theta} \circ \mathbf{m}\), where \(\circ\) denotes the element-wise product. +We can draw a random sample \(\mathbf{m}\) from a Bernoulli distribution \(\text{Bernoulli}(\mathbf{p}=\sigma(\mathbf{x}))\), where \(\mathbf{p}\) is the probability (\(\mathbf{x}\) denotes the logits) of each parameter being activated. However, the discrete Bernoulli distribution is not differentiable, so we cannot backpropagate the gradients through it to optimize the parameters \(\mathbf{p}\) or \(\mathbf{x}\).
+To address this issue, we introduce the Concrete mask which can be drawn from a continuous relaxation of Bernoulli distribution. Before we introduce the Concrete mask, we first review the Gumbel-Max trick in the two-class case.
+Let \(p_0\) and \(p_1\) denote the unnormalized probabilities of a Bernoulli random variable being 0 and 1, respectively, with \(x\) representing the logits. Then, the probability of the event \(m=1\) is given by
+where \(\sigma\) denotes the sigmoid function. +In the context of the Gumbel-Max trick, the occurrence of the event \(m=1\) is determined by the condition \(g_1 + \log p_1 > g_0 + \log p_0\), where \(g_0\) and \(g_1\) are two independent standard Gumbel random variables. +Thus we have
+Because the difference of two standard Gumbel random variables is a Logistic random variable, we can replace \(g_1 - g_0\) by \(\log u - \log(1-u)\) where \(u\) is a random variable sampled from a uniform distribution on the interval \((0,1)\). +Substitute this into Eq.(\ref{eq:appendix_P_m_1}) and express the probability in terms of the logits \(x\) to simplify the expression, we have
+The binary Concrete distribution offers a continuous relaxation of the discrete Bernoulli random variables, which is beneficial for gradient-based optimization as it allows for the backpropagation of gradients even through the sampling process. +Instead of making a hard decision as the above equation, we use a temperature parameter \(\lambda\) to control the steepness of the sigmoid function, and hence control how close our 'soft' decisions are to being 'hard' decisions. The continuous version of the Bernoulli random variable is then given by
+As the temperature \(\lambda\) approaches zero, the sigmoid function becomes a step function, and the Concrete random variable \(\hat{m}\) becomes a Bernoulli random variable, as shown in the following Figure. In the limit when \(\lambda \to 0\), this results in sampling \(m=1\) if \(\log \frac{\sigma(x)}{1 - \sigma(x)} > -\log \frac{u}{1 - u}\), consistent with the original Gumbel-Max trick. +The binary Concrete distribution thus provides a differentiable approximation to Bernoulli random variables. +We can further binarize the Concrete mask by setting the entries with values greater than 0.5 to 1 and the rest to 0.
+
+ 
+Merging CLIP models on eight image classification tasks, using the concrete task arithmetic algorithm
+# tensorboard logs and learned checkpoints of the shared mask can be found at https://huggingface.co/tanganke/clip-vit-base-patch32_concrete-task-arithmetic_tblogs
+fusion_bench \
+ fabric.loggers.name=ViT-B-32/concrete_task_arithmetic \
+ method=concrete_subspace/clip_concrete_task_arithmetic \
+ modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8 \
+ taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
+results
+{
+ "svhn": {
+ "accuracy": 0.903003990650177,
+ "loss": 0.37700024247169495
+ },
+ "stanford_cars": {
+ "accuracy": 0.6326327323913574,
+ "loss": 1.2553859949111938
+ },
+ "resisc45": {
+ "accuracy": 0.7558730244636536,
+ "loss": 1.017554759979248
+ },
+ "eurosat": {
+ "accuracy": 0.9407407641410828,
+ "loss": 0.20871955156326294
+ },
+ "gtsrb": {
+ "accuracy": 0.8285035490989685,
+ "loss": 0.5861473679542542
+ },
+ "mnist": {
+ "accuracy": 0.9800000190734863,
+ "loss": 0.08148527890443802
+ },
+ "dtd": {
+ "accuracy": 0.5249999761581421,
+ "loss": 2.2731478214263916
+ },
+ "sun397": {
+ "accuracy": 0.6421158909797668,
+ "loss": 1.4108904600143433
+ }
+}
+Concrete AdaMerging (Layer-wise)
+# tensorboard logs and learned checkpoints of the shared mask can be found at https://huggingface.co/tanganke/clip-vit-base-patch32_concrete-layer-wise_adamerging_tblogs
+fusion_bench \
+ fabric.loggers.name=ViT-B-32/clip_concrete_layer_wise_adamerging \
+ method=concrete_subspace/clip_concrete_layer_wise_adamerging \
+ modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8 \
+ taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
+
+ X. Yi, S. Zheng, L. Wang, X. Wang, and L. He, “A safety realignment framework via subspace-oriented model fusion for large language models.” arXiv, May 14, 2024. doi: 10.48550/arXiv.2405.09055.
++The paper introduces a safety realignment framework for large language models via subspace-oriented model fusion (SOMF, the authors learn a shared mask on the weight space of large language model), which combines safeguard capabilities of initially aligned models with fine-tuned models to ensure safety without compromising performance on downstream tasks.
+
E. J. Gumbel. Statistical Theory of Extreme Values and Some Practical Applications. A Series of Lectures. Technical +Report PB175818, National Bureau of Standards, Washington, D. C. Applied Mathematics Div., 1954. URL +https://ntrl.ntis.gov/NTRL/dashboard/searchResults/titleDetail/PB175818.xhtml. ↩
+R. Duncan Luce. Individual Choice Behavior. Individual Choice Behavior. John Wiley, Oxford, England, 1959 ↩
+Chris J Maddison, Daniel Tarlow, and Tom Minka. A* sampling. Advances in neural information processing systems, +27, 2014. ↩
+Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The Concrete Distribution: A Continuous Relaxation of Discrete +Random Variables, March 2017. URL http://arxiv.org/abs/1611.00712. ↩
+The DepthUpscalingAlgorithm is used to upscale the depth of PyTorch models. Here's a basic guide on how to use it:
First, import the necessary modules:
+from omegaconf import DictConfig
+from torch import nn
+from fusion_bench.method.depth_upscaling import DepthUpscalingAlgorithm
+from fusion_bench.modelpool import to_modelpool
+Create an instance of DepthUpscalingAlgorithm by passing a configuration dictionary.
+This dictionary should contain the name of the method ("depth_upscaling") and a list of layer indices that determine the upscaling pattern.
method_config = {"name": "depth_upscaling", "layer_indices": [0, 1, 1, 0]}
+algorithm = DepthUpscalingAlgorithm(DictConfig(method_config))
+Assume we have a list of PyTorch models (nn.ModuleList instances) that we want to upscale. Here, we're creating a list of linear models as an example:
Then, we can the model to the run method of our algorithm:
The run method will return an upscaled model. The type of the returned model will be the same as the input models (in this case, nn.ModuleList), and its length will be determined by the layer indices specified in the method configuration.
Here we provide an example of how to use the DepthUpscalingAlgorithm to upscale the depth of a Mistral model 1.
+ from omegaconf import DictConfig
+from torch import nn
+from transformers import AutoModelForCausalLM, MistralConfig, MistralForCausalLM
+from fusion_bench.method.depth_upscaling import DepthUpscalingAlgorithm
+
+# create a Mistral model
+# here we randomly initialize the model for demonstration purposes
+# in practice, you would load a pretrained model
+model_config = MistralConfig(
+ # https://huggingface.co/mistralai/Mistral-7B-v0.1/resolve/main/config.json
+ **{
+ "architectures": ["MistralForCausalLM"],
+ "bos_token_id": 1,
+ "eos_token_id": 2,
+ "hidden_act": "silu",
+ "hidden_size": 4096,
+ "initializer_range": 0.02,
+ "intermediate_size": 14336,
+ "max_position_embeddings": 32768,
+ "model_type": "mistral",
+ "num_attention_heads": 32,
+ "num_hidden_layers": 32,
+ "num_key_value_heads": 8,
+ "rms_norm_eps": 1e-05,
+ "rope_theta": 10000.0,
+ "sliding_window": 4096,
+ "tie_word_embeddings": False,
+ "torch_dtype": "bfloat16",
+ "transformers_version": "4.34.0.dev0",
+ "use_cache": True,
+ "vocab_size": 32000,
+ }
+)
+print('creating model')
+model: MistralForCausalLM = AutoModelForCausalLM.from_config(model_config)
+
+method_config = {
+ "name": "depth_upscaling",
+ "layer_indices": ["range(0,24)", "range(8,32)"],
+}
+algorithm = DepthUpscalingAlgorithm(DictConfig(method_config))
+print('upscaling model')
+upscaled_model = algorithm.run(model.model.layers)
+
+# substitute the model with the upscaled model
+model.model.layers = upscaled_model
+The DepthUpscalingAlgorithm is integrated into the fusion_bench package. You can use it by specifying "depth_upscaling" as the method name in the command line or configuration file.
name: depth_upscaling
+# this should be a list of integers or string, indicating the sequence of layers. If the entry is an integer, it will use the n-th layer of the model. If the entry is a string, it will use the layers specified by the string. The string should be a valid python expression that evaluates to a list of integers.
+# for example, ["range(0,12)", "range(6,12)"] will use the first 12 layers and the last 6 layers of the model to construct the new model
+# [0, 2, 4, "range(6,12)"] will use the 1st, 3rd, 5th, and the 7th to 12th layers of the model to construct the new model
+layer_indices: null
+You can then run the fusion_bench command with the specified configuration file:
DepthUpscalingAlgorithm
+
+
+¶
+ Bases: BaseAlgorithm
Implements the Depth Upscaling Algorithm.
+This class extends the BaseModelFusionAlgorithm to handle depth upscaling of models.
+It supports upscaling the depth of a model by duplicating specified layers.
Parameters:
+layer_indices
+¶list)
+ –
+ List of layer indices to duplicate.
+**kwargs
+¶Additional keyword arguments.
+fusion_bench/method/depth_upscaling/depth_upscaling.pyrun(modelpool)
+
+¶Executes the depth upscaling algorithm on a given model pool.
+This method checks the type of the model pool, ensures that it contains only one model, and verifies that the model is an instance of nn.ModuleList.
Parameters:
+modelpool
+¶ModuleList | ModelPool)
+ –
+ The pool of models to upscale. Must contain only one model.
+Returns:
+ModuleList
+ –
+ nn.ModuleList: The upscaled model.
+Raises:
+AssertionError
+ –
+ If the model pool contains more than one model or if the model is not an instance of nn.ModuleList.
ValueError
+ –
+ If an invalid layer specification is provided in the configuration.
+fusion_bench/method/depth_upscaling/depth_upscaling.pyThe Dummy Algorithm is a simple algorithm that does not perform any fusion operation. Instead, it returns a pretrained model if one is available in the model pool. If no pretrained model is available, it returns the first model in the model pool. +This algorithm is useful for testing and debugging purposes, as it allows you to quickly check if the model pool is set up correctly and the fusion process is working as expected.
+To use the Dummy Algorithm, you need to specify "dummy" as the algorithm name.
The implementation of the Dummy Algorithm is straightforward. Here is the main method of the DummyAlgorithm class:
DummyAlgorithm
+
+
+¶
+ Bases: BaseAlgorithm
fusion_bench/method/dummy.pyThe Fisher merging algorithm 1 is a per-parameter weighed averaging method that assigns weights to the models based on the Fisher information matrix of the models on some labeled data. +The Fisher information matrix \(F_\theta\) of a model with parameters \(\theta\) can be expressed as:
+where \(p(x)\) is the data distribution, \(p(y|x, \theta)\) is the model's output distribution, for example, the softmax output of a classification model, and \(\nabla_\theta\) is the gradient with respect to the model's parameters \(\theta\). +The Fisher information matrix can be used to estimate the importance of each parameter in the model and thus assign weights to the models based on their Fisher information. +In addition, the Fisher information matrix can be used to estimate the similarity between tasks, which can be useful in auxiliary-task learning and multi-task learning scenarios 2.
+As the full Fisher information matrix is often computationally expensive to compute and memory-intensive to store, we approximate using the diagonal Fisher information matrix, which is the diagonal of the full Fisher information matrix. +The diagonal Fisher information matrix can be computed as:
+Assuming we have \(n\) models with parameters \(\theta_i\) and diagonal Fisher information matrices \(\hat{F}_{\theta_i}\), the Fisher merging algorithm computes the merged model's parameters \(\theta\) as follows:
+where \(\theta_i\) are the parameters of the individual models, \(\hat{F}_{\theta_i}\) are the diagonal Fisher information matrices of the individual models, and \(j\) indexes the parameters of the models. +The Fisher merging algorithm can be considered a per-weight weighed averaging method, where the weights are determined by the Fisher information of each parameter in the models.
+Example of merging eight CLIP-ViT-B/32 models using Fisher merging:
+fusion_bench method=clip_fisher_merging \
+ modelpool=clip-vit-base-patch32_TA8 \
+ taskpool=clip-vit-classification_TA8
+Merge eight CLIP-ViT-L/14 models using Fisher merging:
+fusion_bench \
+ method=clip_fisher_merging \
+ method.batch_size=8 method.num_workers=4 \
+ modelpool=clip-vit-large-patch14_TA8 \
+ taskpool=clip-vit-classification_TA8 \
+ taskpool.clip_model=openai/clip-vit-large-patch14
+Merge GPT-2 models for text classification tasks:
+fusion_bench \
+ method=gpt2_fisher_merging \
+ method.num_fisher_examples=512 method.batch_size=8 \
+ modelpool=gpt-2_glue \
+ taskpool=gpt-2_glue
+FisherMergingAlgorithm
+
+
+¶
+ Bases: BaseAlgorithm
Implements the Fisher Merging Algorithm.
+This class extends the BaseModelFusionAlgorithm to handle merging of models using Fisher weights. +It supports excluding certain parameters, normalizing Fisher weights, and setting a minimal value for Fisher weights.
+ + +Methods:
+run
+ –
+ BaseModelPool) -> nn.Module: +Executes the Fisher merging process on the model pool and returns the merged model.
+fusion_bench/method/fisher_merging/fisher_merging.py355 +356 +357 +358 +359 +360 +361 +362 +363 +364 +365 +366 +367 +368 +369 +370 +371 +372 +373 +374 +375 +376 +377 +378 +379 +380 +381 +382 +383 +384 +385 +386 +387 +388 +389 +390 +391 +392 +393 +394 +395 +396 +397 +398 +399 +400 +401 +402 +403 +404 +405 +406 +407 +408 +409 +410 +411 +412 +413 +414 +415 +416 +417 +418 +419 +420 +421 +422 +423 +424 +425 +426 +427 +428 +429 +430 +431 +432 +433 +434 +435 +436 +437 +438 +439 +440 +441 +442 +443 +444 +445 +446 +447 +448 +449 +450 +451 +452 +453 +454 +455 +456 +457 +458 +459 +460 +461 +462 +463 +464 +465 +466 +467 +468 +469 +470 +471 +472 +473 +474 +475 +476 +477 +478 +479 +480 +481 +482 +483 +484 | |
get_fisher_weights(model_name, model, train_dataset, param_names_to_merge)
+
+¶Compute the Fisher weights for the given model and training dataset.
+ + +Parameters:
+model_name
+¶str)
+ –
+ The name of the model.
+model
+¶Module)
+ –
+ The model module.
+train_dataset
+¶The training dataset.
+param_names_to_merge
+¶List[str])
+ –
+ List of parameter names to merge.
+Returns:
+Dict[str, Tensor]
+ –
+ Dict[str, Tensor]: The computed Fisher weights for each parameter.
+fusion_bench/method/fisher_merging/fisher_merging.pyon_fisher_merging_start()
+
+¶Setup the zero-shot classification head before starting the Fisher merging process.
+ +fusion_bench/method/fisher_merging/fisher_merging.pyrun(modelpool)
+
+¶Run the Fisher Merging Algorithm.
+This method constructs the wrapped model and performs test-time adaptation if necessary.
+ + +Parameters:
+modelpool
+¶BaseModelPool)
+ –
+ The model pool containing the pretrained and fine-tuned models.
+Returns:
+Module
+ –
+ nn.Module: The merged model after test-time adaptation.
+fusion_bench/method/fisher_merging/fisher_merging.py388 +389 +390 +391 +392 +393 +394 +395 +396 +397 +398 +399 +400 +401 +402 +403 +404 +405 +406 +407 +408 +409 +410 +411 +412 +413 +414 +415 +416 +417 +418 +419 +420 +421 +422 +423 +424 +425 +426 +427 +428 +429 +430 +431 +432 +433 +434 +435 +436 +437 +438 +439 +440 +441 +442 +443 +444 +445 +446 +447 +448 +449 +450 +451 +452 +453 +454 | |