Trainable Weight Averaging: Accelerating Training and Improving Generalization

Weight averaging is a widely used technique for accelerating training and improving the generalization performance of deep neural networks (DNNs) (Izmailov et al., 2018; Gupta et al., 2019; Yang et al., 2019; Kaddour, 2022; Wortsman et al., 2022). It exploits the linear mode connectivity (Frankle et al., 2020) in DNNs’ loss landscapes, effectively mimicking model ensemble performance within a single weight space. Due to its simplicity and effectiveness, weight averaging has gained significant attention from practitioners.

In SWA (Izmailov et al., 2018), weight averaging is given by $\boldsymbol{w}_{\mathrm{swa}}=\sum_{i=1}^{n}\frac{1}{n}\boldsymbol{w}_{i}$, where $n$ candidate solutions of the network collected at the tail stage of training are equally averaged. This strategy has been shown to be effective in enhancing generalization ability. Other improved strategies include selective averaging methods, such as greedy soup (Wortsman et al., 2022) and latest weight averaging (LAWA) (Kaddour, 2022), as well as leveraging pre-set strategies like exponential moving average (EMA). However, applying equal or pre-set weighting coefficients to candidate model weights can be inappropriate, especially when weights differ significantly, such as those sampled from different training configurations. This can potentially result in suboptimal performance.

In this paper, we study weight averaging with learnable coefficients. Let us regard each model weight as a basis. When we train the coefficients without considering the dependence of different weights, the number of model weights used for averaging is roughly equal to the dimension of the optimization space. It has been reported that for deep models, the training dynamics happen in a low-dimensional subspace (Gur-Ari et al., 2018; Li et al., 2022a) and crucial DNN mode-connectivity patterns emerge early in training (You et al., 2020a; Frankle et al., 2020). Therefore, optimizing the weighting coefficients provides great flexibility and can lead to improvements compared to equal averaging. This is especially important for applying weight averaging in the early stage when the parameters have not been well trained. These findings suggest a promising direction: effectively utilizing these early explorations may enable rapid composition of the final solution while maintaining high accuracy. Since many model weights are involved, optimizing the weighted coefficients also faces practical challenges in both memory and computation. Efficiency is an obstacle that must be solved when dealing with large-scale model training (Wortsman et al., 2022).

In this paper, we delve into efficient and effective weighted averaging of DNNs. We introduce Trainable Weight Averaging (TWA), a novel approach that enables explicit, trainable adjustments to the weighting coefficients. To facilitate efficient optimization, we reframe the problem as a subspace training task: regarding each weight as a point in the full parameter space, we construct a subspace containing the weight points to be averaged. By optimizing in this subspace, we can adaptively search for a good set of weighting coefficients. This scheme, which involves gradient projection onto the subspace, is computationally efficient and can be accelerated by GPUs. To cope with the challenge of serve memory burden when averaging with multiple weights (Wortsman et al., 2022), we develop an efficient distributed subspace training scheme that facilitates parallel processing across multiple nodes. This approach efficiently handles large-scale problems by evenly distributing memory and computational loads across nodes and can be seamlessly integrated with existing distributed training frameworks, such as Distributed Data Parallel (DDP) (Li et al., 2020). Furthermore, to enhance averaging efficiency, we introduce layer-wise processing to better utilize the model’s structure information and find that the matrix associated with gradient projection can be quantized to low-bit representations (e.g., 4 bits) without performance degradation, further reducing memory usage. Consequently, TWA enables efficient subspace training with minimal additional memory overhead, making it scalable for averaging large models.

By exploiting the low-dimensional nature of the subspace, we reduce the degrees of freedom from millions or billions to just dozens or hundreds. This reduction allows us to train with a relatively small number of samples. We, therefore, propose optimizing these coefficients using a small held-out validation set, leading to our variant TWA-v, in contrast to the version using training data, TWA-t. Since validation data are unseen in the training set, they could serve as new data to evaluate generalization capability and have played a crucial role in machine learning. For example, we use the validation set to assess the model checkpoints during the training and select the best one. We now have new ways to better utilize these model checkpoints by appropriately averaging with TWA-v, resulting in improved overall performance.

We demonstrate TWA’s effectiveness in two scenarios through extensive experiments across model architectures and tasks: 1) accelerating training by averaging historical solutions during the head stage of training, and 2) improving generalization by averaging solutions fine-tuned from single or multiple fine-tuning configurations. Our results show over 40% and 30% reduction in training epochs on CIFAR-100 and ImageNet respectively, while maintaining or improving performance (Tables 4, 5). We also demonstrate that TWA can better utilize the fine-tuned models, reducing the number of models required by 4x for reaching comparable performance compared to previous methods (Table 16). Interestingly, we also find that TWA-v exhibits particular effectiveness with transformer architectures.

In summary, our contributions can be summarized as follows:

1. 1.

   We propose Trainable Weight Averaging (TWA), an efficient approach that enables the averaging of weights with layer-wise learnable coefficients to accelerate training and improve the generalization of DNNs.

2. 2.

   We design an efficient scheme to handle large-scale problems via subspace training, enabling multi-node parallel training by evenly distributing the memory and computation burden into different nodes. We also devise a compression strategy to further reduce the memory footprint, facilitating effective coefficient optimization.

3. 3.

   Based on the fact that the number of learnable variables is small, we propose TWA-v that optimizes the coefficients with a small held-out validation set to enable more efficient and effective averaging. Additionally, we discover that TWA-v is particularly effective for transformer-based architectures.

4. 4.

   We demonstrate the efficiency and effectiveness of TWA through extensive experiments involving various architectures (e.g., CNNs, ViTs, and GPT-2), different tasks (e.g., image classification, machine translation, and language modeling), and multiple training scenarios (from scratch training to fine-tuning) using different optimizers (SGD and Adam).

Training neural networks in subspaces has recently become an interesting topic, garnering considerable interest from researchers (Vinyals and Povey, 2012; Gur-Ari et al., 2018; Tuddenham et al., 2020). The pioneering work (Li et al., 2018) first proposed training neural networks in a reduced random subspace to measure the intrinsic dimension of loss objectives. The following work (Gressmann et al., 2020) improved the training performance of random bases by considering the layer-wise structure and re-drawing the random bases at each step. Instead of utilizing random bases, Li et al. (2022a) proposed a low-dimensional trajectory hypothesis and extracted the subspaces from historical training dynamics, dramatically improving the dimensionality efficiency. Then Li et al. (2022b) applied subspace training to adversarial training, effectively addressing the existing catastrophic and robust overfittings, and thereby significantly improving the model robustness performance. In this paper, we reframe trainable weight averaging as a subspace training problem and develop an efficient training scheme that enables multi-node parallel training for large-scale problems, incorporating improved subspace construction procedures.

A lot of efforts have been made to speed up DNNs’ training (Shen et al., 2023). Apart from the well-known optimization methods on adaptive learning rates, e.g., Adam (Kingma and Ba, 2015) or accelerated schemes, e.g., Nesterov momentum (Nesterov, 1983), Zhang et al. (2019) proposed LookAhead optimizer that utilizes the search direction generated by another “fast” optimizer, achieving faster convergence and better learning stability. Goyal et al. (2017) adapted a large mini-batch to speed up the training and introduced a scaling rule for adjusting the learning rates. You et al. (2017, 2020b) proposed a layer-wise adaptive learning rate to further scale the batch size and shorten the training time. Gupta et al. (2019) proposed to use large mini-batches to compute an approximate solution quickly and then refined it by averaging the weights of multiple models computed independently and in parallel to accelerate the training. In this paper, we improve the DNNs’ training efficiency by sufficiently utilizing the historical solutions during the training and conducting training in a subspace with substantially reduced dimensions. In this way, we significantly reduce the required training epochs.

Next, we review previous weight-averaging strategies.

In this paper, we propose trainable weight averaging (TWA), a method that optimizes the weighting coefficients of different model weights to improve performance. Given a set of weights $\left\{\boldsymbol{w}_{i}\right\}_{i=1}^{n}$, which can be sampled either from a single training trajectory or collected from multiple fine-tuning configurations, we aim to linearly combine them for creating an improved averaged solution. Specifically, the set of possible TWA solutions we considered, i.e., $\boldsymbol{w}_{\mathrm{twa}}$, can be represented as follows:

One can see that the solution set $U$ forms a linear space with dimension $n$.

Previous pre-set averaging strategies may fall short for complex weights, potentially leading to averaged solutions that are suboptimal within the subspace $U$. In fact, the solutions obtained by these approaches can be viewed as specific points within $U$. TWA addresses this limitation by optimizing the weighting coefficients to achieve better performance. A comparison between TWA and previous averaging approaches is presented in Table 2.

We now consider how to search for the optimal solution $\boldsymbol{w}^{*}_{\mathrm{twa}}$ within the subspace ${U}$. First, we will identify a set of bases $\{\boldsymbol{e}_{i}\}_{i=1}^{n}$ to support the solution space such that ${U}=\left\{\beta_{1}\boldsymbol{e}_{1}+\beta_{2}\boldsymbol{e}_{2}+\cdots+\beta_{n}\boldsymbol{e}_{n}\leavevmode\nobreak\ \leavevmode\nobreak\ |\leavevmode\nobreak\ \leavevmode\nobreak\ \beta_{i}\in\mathbb{R}\right\}$, where $\{\beta_{i}\}_{i=1}^{n}$ are the coefficients to be optimized. Rather than using $\{{\boldsymbol{w}_{i}}\}_{i=1}^{n}$ directly as bases, we would decouple them to $\{\boldsymbol{e}_{i}\}_{i=1}^{n}$ since the weight vectors $\{{\boldsymbol{w}_{i}}\}_{i=1}^{n}$ are often highly correlated, particularly when fine-tuned from the same pre-trained model. Such correlations complicate optimization and can degrade the training performance, as shown in Section 5.3.4. Then we search for $\boldsymbol{w}^{*}_{\mathrm{twa}}$ in ${A}$ by optimizing $\beta_{i}$ with the following problem,

The second term serves as a regularization for $\beta_{i}$ with a coefficient $\lambda>0$.

There are two choices for using the available data to optimize Eqn. (5):

• •

  Using the training data, i.e., $\boldsymbol{S}=\boldsymbol{S}^{\text{train}}$, while still leveraging the validation data $\boldsymbol{S}^{\text{val}}$ to select the best-performing model during training after each training epoch.

• •

  Using the validation data directly, i.e., $\boldsymbol{S}=\boldsymbol{S}^{\text{val}}$, to optimize the weighting coefficients and select the final model after the training is complete.

The latter approach can be more efficient and effective: the rationale behind this choice is that since we only need to optimize a small number of coefficients, a small held-out validation set $\boldsymbol{S}^{\text{val}}$ is sufficient. Given that the weights $\{\boldsymbol{w}_{i}\}_{i=1}^{n}$ are unseen to $\boldsymbol{S}^{\text{val}}$, if $\boldsymbol{w}_{\rm twa}$ performs well on the validation set, it is expected to generalize effectively during testing. Additionally, the validation set can be typically much smaller than the training set, which helps reduce training costs. We denote this validation-based approach as TWA-v, distinguishing it from the training-based approach denoted as TWA-t. This introduces a novel use of validation data by learning optimal weighting coefficients to combine multiple models, rather than selecting a single best-performing model.

A straightforward approach to optimizing the coefficients $\{\beta_{i}\}_{i=1}^{n}$ is to construct a computational graph that links the bases $\left\{\boldsymbol{e}_{i}\right\}_{i=1}^{n}$ and coefficients to the model weights $\boldsymbol{w}$, enabling gradient descent on $\{\beta_{i}\}_{i=1}^{n}$ through standard backpropagation in an end-to-end manner. However, as noted by Wortsman et al. (2022), this approach results in a large computational graph, introducing excessive additional memory overhead and becoming prohibitive as the model size and number of weights increase.

In this paper, we propose a simpler and more efficient method for optimizing the coefficients $\beta_{i}$ without constructing an additional computational graph, by leveraging subspace training. We notice that there exists a mapping between the coefficient space $\{\beta_{i}\}_{i=1}^{n}\in\mathbb{R}^{n}$ and the parameter space $\mathbb{R}^{D}$:

This implies that each set of coefficients $\{\beta_{i}\}_{i=1}^{n}$ uniquely corresponds to a point $\boldsymbol{w}$ in the parameter space $\mathbb{R}^{D}$, thereby forming a subspace of dimensionality $n$. As a result, we can optimize the coefficients $\{\beta_{i}\}_{i=1}^{n}$ by optimizing $\boldsymbol{w}$ within this subspace, which is more efficient than performing direct optimization through a computational graph.

We then present our subspace training scheme for efficient trainable weight averaging, which consists of subspace construction and training. To further enhance performance and efficiency, we introduce a layer-wise extension to incorporate model structure, develop a multi-node distributed training approach for large-scale problems, and quantize the projection matrix to reduce memory overhead. Before detailing the training procedure, we first discuss the challenges of averaging multiple model weights.

• •

  Computation. As the model size and the number of models grow, the costs of both subspace construction and subspace training increase accordingly.

• •

  Memory. Subspace training requires storing a matrix for projecting gradients onto the subspace, which can be challenging in large-scale scenarios where the matrix may exceed a single GPU’s memory capacity.

• •

  Balance load. Both subspace construction and training are best performed on GPUs, with the computational load distributed across multiple nodes when available. This significantly reduces training time and improves the scalability of the algorithm.

In this section, we conduct numerical experiments on various computer vision and natural language processing tasks to demonstrate the efficiency and effectiveness of our proposed TWA approach. First, we show that TWA can significantly accelerate the training of DNNs by averaging historical solutions from the head stage of training, especially for transformer-based architectures. Next, we demonstrate that TWA significantly improves the performance of fine-tuned models across both single and multiple training configurations. Finally, we present ablation studies to further analyze the properties of TWA training.

In this work, we propose Trainable Weight Averaging (TWA), an efficient framework that enables weight averaging with learnable weighting coefficients. It extends the manually defined weighting coefficients as in previous works to a trainable manner, which endows with much greater flexibility and enables handling weights from different stages and configurations. We design an efficient parallel training scheme to cope with large-scale training and propose quantization schemes for the projection matrix to achieve memory efficiency. Additionally, we derive two variants, TWA-t and TWA-v, based on the data used for training, and show that TWA-v allows for more efficient and effective averaging when validation data is available. Extensive experiments on both efficient training and fine-tuning tasks demonstrate the effectiveness and efficiency of our approach.