Geometry-Aware Collaborative Multi-Solutions Optimizer for Model Fine-Tuning with Parameter Efficiency

Dear reviewer,

We thank the reviewer for your valuable the time to review our work. Below, we address the point you have raised.

Training and inference time When $M > 1$, the training process involves the loss function $l_{div}$, which requires calculating the determinant of the prediction matrix, and the kernel function $K(., .)$, which involves calculating distances in the model or latent space. As a result, increasing M leads to an exponential increase in training time. However, during inference, since the loss and kernel functions are not involved, the shared architecture among particles can be done in parallel. Therefore, increasing M does not significantly impact the inference time. In practice, the number of particles is selected based on the available system resources and the expected performance. We consider the trade-off between performance and training time when increasing the number of solutions M as our limitation.

Table 1: Training and inference time (s) on 1 epoch

We agree that increasing the number of particles leads to a linear increase in training time. However, this is a common characteristic shared by particle-based sampling methods (e.g., SVGD and SGLD) as well as ensemble approaches.

We kindly ask the reviewer to consider the theoretical and technical contributions of our work. Our study focuses on particle-based sampling and variational gradient approaches, including well-established methods such as Stochastic Gradient Langevin Dynamics (SGLD) and Stein Variational Gradient Descent (SVGD), which are widely used in Bayesian inference (e.g., Bayesian Neural Networks). Specifically, we consider the gradient flow in the probability space that connects the prior and posterior distributions. We develop theoretical foundations to endow this gradient flow with a well-defined geometry and introduce a divergence term that encourages collaborative yet diverse particle solutions. In our experiments, we compare our method against relevant baselines in Bayesian inference and particle sampling, demonstrating its significant improvement over state-of-the-art methods.

Dear Reviewer,

Thanks for your reminder.

We have mentioned the limitation of our work when discussing training and inference time in the rebuttal, where we stated: "We consider the trade-off between performance and training time when increasing the number of solutions (i.e., $M$) as the limitation of our work." This limitation is also a common characteristic shared by particle-based sampling methods as well as ensemble approaches. We will carefully discuss this limitation in the revised version.

We sincerely thank the reviewer for your time to review our work. Below, we provide a response to the questions you raised.

Q1: The importance of geometry

Geometry plays a crucial role in our approach. It is derived from the proximal operator in Eq. (3), which incorporates the Fisher information as defined in Eq. (4). From this perspective, the resulting term acts as a geometry-aware regularization. Furthermore, as shown in Eq. (4), the geometry influences the update along each dimension of the model parameters.

Q2: Training and inference time When $M > 1$, the training process involves the loss function $l_{div}$, which requires calculating the determinant of the prediction matrix, and the kernel function $K(., .)$, which involves calculating distances in the model or latent space. As a result, increasing M leads to an exponential increase in training time. However, during inference, since the loss and kernel functions are not involved, the shared architecture among particles can be done in parallel. Therefore, increasing M does not significantly impact the inference time. In practice, the number of particles is selected based on the available system resources and the expected performance. We consider the trade-off between performance and training time when increasing the number of solutions M as our limitation.

Table 1: Training and inference time (s) on 1 epoch

Q3: Training algorithm Here we present the training algorithm of our GAC-MSO. In our algorithm, we need to compute the loss function $\psi(.)$, $l_{div}(.)$, the kernel function $K(.,.)$, and their respective gradients. In practice, the loss function $\psi(.)$ (which we use as the cross-entropy loss) and $l_{div}(.)$ are calculated using the output logits, and the gradients of these functions are derived simultaneously using the standard backpropagation algorithm. For the kernel $K$, we employ the RBF kernel, which is applied to the model space. The RBF function uses the $l_2$ distance, a computationally simple and efficient measure with a closed-form for its gradient. However, we agree that the overall process adds to the computational cost of the training.

Input: Initialize particles $\theta_1, \theta_2, ..., \theta_M$, kernel $K(., .)$, momentum $\gamma=0.9$, trade-off $\beta=1$, $\alpha=0.2$, and learning rate $\eta$

For each particle $i=1...M$:

do

[1]. Compute the empirical estimate of the score function:

$h\left(\theta_i\right)= \sum_{m=1}^{M}\left[\beta\nabla_{\theta_m}\Psi\left(\theta_m\right)K\left(\theta_m,\theta_i\right)-\frac{\sum_{m=1}^{M}\left[K\left(\theta_m,\theta_i\right)\nabla_{\theta_m} K\left(\theta_m,\theta_i\right)\right]}{\sum_{m=1}^{M}K\left(\theta_m,\theta_i\right)}\right]+\alpha\sum_{m=1}^{M}\nabla_{\theta_{m}}l_{\text{div}}\left(\theta_{1:M}\right)K\left(\theta_{m},\theta_i\right).$

$H(\theta_i) = \left(\nabla_{\theta_m}\Psi\left(\theta_m\right) + \nabla_{\theta_{m}}l_{\text{div}}\left(\theta_{1:M}\right)\right)^2$

[3] Using a moving average with momentum:

$m(\theta_i) = \gamma H(\theta_i) + (1-\gamma)m(\theta_i)$

[4] Update the particle using the gradient descent step:

$\theta_i = \theta_i - \frac{\eta}{M(m(\theta_i) + \epsilon)}h(\theta_i)$

done

Output: Return the updated particles $\theta_1, \theta_2, ..., \theta_M$.

Q4: The choice of hyperparameters In this paper, we formulate the problem in a general form, so that multiple hyperparameters are used, such as the number of particles $M$, the trade-off parameters $\alpha$ and $\beta$ for the loss functions, and the kernel function $K(., .)$. However, in the main experiments, we set $M = 4$, $\alpha = 1$, $\beta = 0.2$, and use an RBF kernel for $K(., .)$ to focus on analyzing the effectiveness of our algorithm in learning diverse solutions with geometry collaboration. The sensitivity of $M$ and $\beta$ is explored in the supplementary material, which demonstrates the effectiveness of each hyperparameter. For the kernel function $K(.,.)$, we use the RBF kernel, as it is a common function to measure the similarity between two points in space. This simple kernel is effective for high-dimensional spaces, such as model or latent spaces.

Dear Reviewer,

We sincerely thank you for taking the time to review our work. Below, we provide a response to the questions you raised.

Q1: Idea from [8] and [25]

Thanks for this comment. The paper [8] introduces a general family of a Blob approach for diffusion processes. In our paper, we leverage this Blob approach to our framework. Moreover, drawing inspiration from Determinantal Point Processes (DPP) theory [25], we develop our divergence term.

Q2: Training and inference time When $M > 1$, the training process involves the loss function $l_{div}$, which requires calculating the determinant of the prediction matrix, and the kernel function $K(., .)$, which involves calculating distances in the model or latent space. As a result, increasing M leads to an exponential increase in training time. However, during inference, since the loss and kernel functions are not involved, the shared architecture among particles can be done in parallel. Therefore, increasing M does not significantly impact the inference time. In practice, the number of particles is selected based on the available system resources and the expected performance. We consider the trade-off between performance and training time when increasing the number of solutions M as our limitation.

Table 1: Training and inference time (s) on 1 epoch

We cannot compare to [8,25] because they are only the building-blocks to develop our approach. [8] proposes a general family of a Blob approach for general diffusion processes. We use that in Eq. (11) in the second term ($K*\rho$) to make smooth the entropy function and come up with the update in Theorem 4.3. Furthermore, [25] discusses how to use determinantal point processes for machine learning in general which we base on to encourage the diversity of the model predictions.

We appreciate the time and effort you put into reviewing our paper. We truly appreciate your thoughtful feedback and valuable comments, and we have addressed the points you raised below.

Q1: Comparison with ensemble of multiple PETF methods

We thank the reviewer for pointing out the interesting paper [1]. However, we believe that it does not compromise the novelty and technical contributions of our work. Specifically, our study focuses on particle-based sampling or variational gradient approaches, including well-established methods such as Stochastic Gradient Langevin Dynamics (SGLD) and Stein Variational Gradient Descent (SVGD), which are widely used in Bayesian inference (e.g., Bayesian Neural Networks). Technically, we consider the gradient flow in the probability space that connects the prior and posterior distributions. We then develop theoretical foundations to endow this gradient flow with a geometry, along with a divergence term that promotes collaborative yet diverse particle solutions. In our experiments, we compare our method against relevant baselines in Bayesian inference and particle sampling.

Motivated by your suggestion, we experimented using a deep ensemble of four different PEFT components following [1]. For our approach, inspired by [2], we compute the kernel based on model representations—that is, we feed the current mini-batch into four particle models and use their representations to compute the kernel similarity. It is worth noting that in the original results reported in the paper, the kernel was computed based on particle parameters.

Table 1: Comparison with ensemble of multiple PETF methods with GAC-MSO

The experimental results show that our approach outperforms the ensemble of four distinct particles, demonstrating the effectiveness of our proposed collaborative and divergent solution.

We also acknowledge the strength of the ensemble approach with four different PEFT particles. This simple strategy provides sufficiently diverse solutions that enhance accuracy. This observation motivates a compelling research question: how can we design a new particle sampling approach in which each particle employs a distinct PEFT type? We leave this promising direction for future work.

[2] Xingchao Liu, Xin Tong, and Qiang Liu. Profiling pareto front with multi-objective stein variational gradient descent. NeurIPS 2021.

Q2: Output ensembling

For the final output, we average the logits before they are passed into the SoftMax function to calculate the accuracy, which is a standard ensembling approach. This will be clearly explained in our revised version.

Q3: Training and inference time

When $M > 1$, the training process involves the loss function $l_{div}$, which requires calculating the determinant of the prediction matrix, and the kernel function $K(., .)$, which involves calculating distances in the model or latent space. As a result, increasing M leads to an exponential increase in training time. However, during inference, since the loss and kernel functions are not involved, the shared architecture among particles can be done in parallel. Therefore, increasing M does not significantly impact the inference time. In practice, the number of particles is selected based on the available system resources and the expected performance. We consider the trade-off between performance and training time when increasing the number of solutions M as our limitation.

Table 2: Training and inference time (s) on 1 epoch

Q4: Training algorithm

Due to NeurIPS policy, we are unable to use external links for figures. Instead, we present the training algorithm of our GAC-MSO here for a clearer understanding.

Input: Initialize particles $\theta_1, \theta_2, ..., \theta_M$, kernel $K(., .)$, momentum $\gamma=0.9$, trade-off $\beta=1$, $\alpha=0.2$, and learning rate $\eta$

For each particle $i=1...M$:

do

[1]. Compute the empirical estimate of the score function:

$h\left(\theta_i\right)= \sum_{m=1}^{M}\left[\beta\nabla_{\theta_m}\Psi\left(\theta_m\right)K\left(\theta_m,\theta_i\right)-\frac{\sum_{m=1}^{M}\left[K\left(\theta_m,\theta_i\right)\nabla_{\theta_m} K\left(\theta_m,\theta_i\right)\right]}{\sum_{m=1}^{M}K\left(\theta_m,\theta_i\right)}\right]+\alpha\sum_{m=1}^{M}\nabla_{\theta_{m}}l_{\text{div}}\left(\theta_{1:M}\right)K\left(\theta_{m},\theta_i\right).$

$H(\theta_i) = \left(\nabla_{\theta_m}\Psi\left(\theta_m\right) + \nabla_{\theta_{m}}l_{\text{div}}\left(\theta_{1:M}\right)\right)^2$

[3] Using a moving average with momentum:

$m(\theta_i) = \gamma H(\theta_i) + (1-\gamma)m(\theta_i)$

[4] Update the particle using the gradient descent step:

$\theta_i = \theta_i - \frac{\eta}{M(m(\theta_i) + \epsilon)}h(\theta_i)$

done

Output: Return the updated particles $\theta_1, \theta_2, ..., \theta_M$.

Citation confusion The result of the FFT is taken from paper VPT. And the experiments for single solution using AdamW are taken from the LoRA paper because in this paper, the authors used AdamW to train. We think it will be fair to use the original results.

Thank you for your comments and suggestions. In our paper, we report ensemble accuracies and Expected Calibration Error (ECE), which are widely used in prior works on particle sampling and Bayesian Neural Networks. Inspired by your suggestions, we further investigate the robustness of our approach using the relevant methodology introduced in [1].

Table 3. Comparison between GAC-MSO and WiSE. For WiSE, the number of alpha present for the number of models to ensemble

In Table 3, we compare our GAC-MSO (using 4 particles with 4 LoRA models) to both Deep-ensemble (also using 4 particles with 4 LoRA models) and Merge-space WiSE (which trains a single LoRA but ensembles multiple models using different alpha values), as requested by the reviewer. The number of alphas in WiSE corresponds to the number of ensemble models, with the exact alpha values for each setting as follows: 4 alphas using [0.1, 0.4, 0.7, 1]; 5 alphas using [0.1, 0.3, 0.5, 0.7, 0.9]; and 10 alphas using [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]. Our results show that, despite WiSE's implicit flattening of the loss landscape (as the original proposal for full fine-tuning), the outputs from different alpha values lack sufficient diversity to improve robustness when compared to Deep-ensemble and GAC-MSO.

Dear reviewer,

As the discussion period is coming to an end, we would like to express our appreciation to the reviewer for reviewing our paper and pointing out such interesting work. We would appreciate it if the reviewer could let us know if there are any questions left unresolved, so we can address them promptly before finalizing our rebuttal.