KIND: Knowledge Integration and Diversion in Diffusion Models

FSGAN (Robb et al. 2020) introduces spectral shifts by directly applying SVD to pre-trained model parameters and fine-tune singular values for model adaptation. Similar approaches are used by (Sun et al. 2022) and (Han et al. 2023) for segmentation and image generation, respectively. The success of these methods demonstrates that SVD can create a compact parameter space in pre-trained models, facilitaing efficient fine-tuning.

However, directly applying SVD to pre-trained parameter matrices will decompose them according to fixed orthogonalization rules, which reduces the interpretability of the singular vectors and makes it challenging to determine which singular vectors contain common knowledge that is suitable for transfer. To address this issue, we employ knowledge integration by reconstructing weight matrices using the SVD-derived matrix forms $U$, $\Sigma$ and $V$, rather than directly applying SVD to the pre-trained parameter matrices.

For the DiT architecture, the main weight matrices in a $L$-layer DiT are $\mathcal{W}=\{W_{q}^{(1\thicksim L)},W_{k}^{(1\thicksim L)},W_{v}^{(1\thicksim L)},W_{o}^{(1\thicksim L)},W_{in}^{(1\thicksim L)},W_{out}^{(1\thicksim L)}\}$. Let $W_{\star}^{(l)}$ represent any weight matrix in layer $l$, where $\star\in\mathcal{S}$ and $\mathcal{S}=\{q,k,v,o,in,out\}$ denotes the set of subscripts. The matrices $U_{\star}^{(l)}$, $\Sigma_{\star}^{(l)}$, $V_{\star}^{(l)}$ are the corresponding components that constitute $W_{\star}^{(l)}$, which is calculated as:

where $\Sigma=\text{diag}(\boldsymbol{\sigma})$ with $\boldsymbol{\sigma}=[\sigma_{1},\sigma_{2},...,\sigma_{r}]$. Here, $W_{\star}^{(l)}\in\mathbb{R}^{m_{1}\times m_{2}}$, $U_{\star}^{(l)}=[u_{1},u_{2},...,u_{r}]\in\mathbb{R}^{m_{1}\times r}$, $\Sigma_{\star}^{(l)}\in\mathbb{R}^{r\times r}$ and $V_{\star}^{(l)}=[v_{1},v_{2},...,v_{r}]^{\top}\in\mathbb{R}^{r\times m_{2}}$. The rank $r$ and dimensions $m_{1}$ and $m_{2}$ define the sizes of $W_{\star}^{(l)}$. By updating $U_{\star}^{(l)}$, $\Sigma_{\star}^{(l)}$ and $V_{\star}^{(l)}$, the $W_{\star}^{(l)}$ can be updated accordingly.

Given a dataset $\mathcal{D}$ with $N_{cls}$ categories, our objective is to diverse knowledge during the training of DiTs. We categorize the components of $U$, $\Sigma$, and $V$ (i.e., row/column vectors $u_{i}$, $\sigma_{i}$, and $v_{i}$) that condense common knowledge as learngenes, while those representing class-specific knowledge are termed tailors. Specifically, we partition the components in $U$, $\Sigma$, and $V$ based on the number of categories $N_{cls}$ and the matrix rank $r$, which satisfies $r=N_{G}+N_{cls}\cdot N_{T}$, where $N_{G}$ is the number of components condensing common knowledge, which make up the learngenes $\mathcal{G}=\{\mathcal{G}_{\star}^{(l)}|\star\in\mathcal{S}\;\text{and}\;l\in[1,L]\}$, with:

Here $U_{G}=\{u_{1},u_{2},\ldots,u_{N_{G}}\}$ ($U_{G}=u_{1\thicksim N_{G}}$ for shot), $\Sigma_{G}=\sigma_{1\thicksim N_{G}}$, and $V_{G}=v_{1\thicksim N_{G}}$.

The $N_{T}$ represents the number of components corresponding to each category, forming the tailor $\mathcal{T}=\{\mathcal{T}_{i,\star}^{(l)}|\star\in\mathcal{S},\,l\in[1,L]\;\text{and}\;i\in[1,N_{cls}]\}$, with:

where $U_{T_{i}}=u_{(t_{i}\thicksim t_{i}+i\cdot N_{t})}$, $\Sigma_{T_{i}}=\sigma_{(t_{i}\thicksim t_{i}+i\cdot N_{t})}$ and $V_{T_{i}}=v_{(t_{i}\thicksim t_{i}+i\cdot N_{t})}$. Here, $t=N_{G}+i\cdot N_{T}$. Thus, each weight matrix in $U$, $\Sigma$, and $V$ is decomposed into a learngene $\mathcal{G}$ and $N_{cls}$ tailors $\mathcal{T}$.

During the training of DiTs, we introduce a class gate $G=[0,\ldots,0,1,0,\ldots,0]\in\mathbb{R}^{N_{cls}}$, where only one element is set to 1, corresponding to the class index. This mechanism ensures that for each training class, only the weight parameters of the learngenes and relevant tailors are updated, facilitating targeted knowledge diversion. The optimization objective is defined as:

where the loss function $\mathcal{L}$ is as defined in Eq.1.

After diverting the knowledge, we can obtain the learngenes and tailors, which condense common knowledge and class-specific knowledge, respectively. Thus, when transferring pre-trained models to novel tasks, only learngenes need to be transferred, significantly improving transfer efficiency. Then, based on the difficulty of downstream tasks (i.e., the number of classes), the corresponding number of tailors are randomly initialized and concatenated with learngenes to form the weight matrices $\hat{W}_{\star}^{(l)}$ of descendants models as:

where $\hat{U}_{T_{1}\thicksim T_{n}}$, $\hat{\Sigma}_{T_{1}\thicksim T_{n}}$, and $\hat{V}_{T_{1}\thicksim T_{n}}$ are composed of $T_{n}$ randomly initialized tailors built for specific tasks.

During fine-tuning, we freeze the parameters of learngenes and only update tailors, allowing them to learn class-specific knowledge from downstream tasks, thereby achieving more efficient fine-tuning.

To better integrate and divert knowledge, we conduct experiments on ImageNet-1K, which comprises 1,000 classes, with 1.2M training images and 50K validation images. To minimize inter-class similarity, we further merge certain similar classes based on the superclasses in WordNet (Miller 1995), resulting in a final total of 611 classes. Among them, 150 classes are used for pre-training diffusion models and condensing learngenes, and while the remaining 461 classes serve as novel classes for evaluating the performance of learngenes and other PEFT methods. Further details are provided in Appendix.

For pre-training DiT and extracting learngenes, we train class-conditional latent DiTs of sizes -B and -L, with a latent patch size of $p=2$ at a $256\times 256$ image resolution on training classes. All models are trained using AdamW with a batch size of 256 and a constant learning rate of $1\times 10^{-4}$ over 200K steps. When fine-tuning on novel classes, we randomly divide 461 novel classes into 18 image generation tasks, each generating images for $c$ classes, where $c\in[7,35]$. All PEFT and learngene methods are fine-tuned with a constant learning rate of $1\times 10^{-3}$ over 50K steps. We use an exponential moving average (EMA) of DiT weights with a decay rate of 0.9999, and results are reported using the EMA model. During image generation, a classifier-free guidance (cfg) scale of 1.5 is applied. Performance is evaluated using Fréchet Inception Distance (Heusel et al. 2017), sFID (Nash et al. 2021), Inception Score (Salimans et al. 2016) and Precision/Recall (Kynkäänniemi et al. 2019).

We first compare KIND with several PEFT methods, which are 1) OFT (Qiu et al. 2023): Multiplies pre-trained weights by a trainable orthogonal matrix. 2) LoRA (Hu et al. 2022): Replaces adapters (Hu et al. 2023) with two low-rank matrices to further reduce parameters. 3) PiSSA (Meng, Wang, and Zhang 2024): Fine-tunes only principal components of the original matrix based on LoRA. 4) SVDiff (Han et al. 2023): Applies SVD to pre-trained weight matrices and fine-tunes only the singular values. Additionally, we adapt two learngene methods for DiTs. 5) Heur-LG (Wang et al. 2022): Selects layers to be transferred based on gradients during continual learning. 6) Auto-LG (Wang et al. 2023): Uses meta-networks to select layers that are similar to downstream tasks. Lastly, we provide the results for full fine-tuning. 7) Full FT: Directly fine-tuning all parameters of pre-trained models.

To evaluate the adaptability of KIND, we compare it with several PEFT and learngene methods on novel tasks fairly, ensuring that the pre-trained models and learngenes used in these methods are trained with the same setting. As shown in Table 1, our proposed KIND achieves state-of-the-art results on DiT-B/2 and DiT-L/2, demonstrating significant improvements (in DiT-L/2) with more than 6.54 and 1.07 decrease in FID and sFID respectively with only 45.4M trainable parameters and the computational cost saved at least 35.4G FLOPs. The improvements in IS ($\uparrow$38.8) and precision ($\uparrow$0.02) further underscore the superiority of KIND.

We observe a clear performance gap between all PEFT methods and Full FT, despite the computational efficiency and reduced trainable parameters offered by PEFT methods. This disparity underscores a substantial difference between the novel tasks and the training tasks. Consequently, directly fixing the parameters of pre-trained models, as done in PEFT methods, might not be an optimal strategy in such scenarios. As illustrated in Figure 5, the images generated by PEFT methods fail to adequately capture the knowledge of corresponding categories due to the limited number of trainable parameters and the significant gap between the knowledge in pre-trained models and the target tasks.

Existing learngene methods, such as Heur-LG and Auto-LG, have not achieved the same level of success in image generation tasks as they have in image classification (Wang et al. 2022, 2023). These methods selectively transfer knowledge in pre-trained models while retaining some flexibility to acquire knowledge from novel tasks. However, they still heavily rely on pre-trained models trained with traditional objectives and do not intervene in the process of knowledge acquisition of pre-trained models, leading to sub-optimal results. Additionally, these methods introduce too many randomly initialized parameters. While this can be beneficial for parameter transfer, it hinders the effective learning of useful knowledge from a limited number of images.

Conversely, our KIND shifts the learning objectives from simply maximizing model performance on training dataset to condensing as much transferable common knowledge as possible. KIND emphasizes the extraction of highly transferable knowledge during the pre-training stage, using knowledge diversion to condense common knowledge into learngenes, thereby making learngenes a stronger backbone for task adaptability. Additionally, we apply low-rank assumptions to tailors, making them class-specific, with their rank flexibly set based on task difficulty, further ensuring structural adaptability. As shown in Figure 5 and Tabel 1, the images generated by KIND are significantly better than those produced by other PEFT methods, both visually and in terms of performance metrics. Surprisingly, KIND’s performance even surpass Full FT with significant visual improvements while saving 411.4M trainable parameters and reducing computational cost by 35.4G FLOPs.

As noted in (Wang et al. 2022; Xia et al. 2024), learngenes can accelerate the adaptation of descendant models to novel tasks by transferring common knowledge, providing a significant advantage over training from scratch. Beyond this, our proposed KIND further enhances the convergence speed compared to PEFT methods. Figure 4 illustrates the convergence speed of KIND and other PEFT methods, showcasing images generated by the models at every 10K training steps.

The convergence speed is typically influenced by the number of trainable parameters for fine-tuning, with PEFT methods primary aiming to reduce this number through techniques like orthogonalization and low-rank constraints. However, these methods often overlook the importance of the transferability of knowledge in pre-trained models, as they tend to fix the pre-training parameters of backbone network. In contrast, KIND uses the learngenes that condense common knowledge as the backbone, providing superior transferability while maintaining lightweight. Meanwhile, the tailors ensure the models acquire task-specific knowledge, allowing KIND to achieve faster convergence and better performance on downstream tasks.

We validate the effectiveness of learngenes and tailors along with class gate through ablation experiments. #1 performs SVD on pre-trained weights and randomly selects $N_{G}$ singular vectors comprising its backbone and then fine-tune it with LoRA. #2 uses the learngenes extracted by KIND as the backbone based on #1. #3 uses tailors instead of LoRA to fine-tune models without class gate.

As shown in Table 2, the knowledge in the learngenes which have undergone knowledge diversion, is more common and thus better suited for adapting to downstream tasks, especially when these tasks differ significantly from the training tasks. Futhermore, the tailors themselves function as a PEFT method by utilizing a low-rank assumption to combine class-specific knowledge into the pre-trained models or learngenes, effectively augmenting the backbone network with additional components. This combination helps the model better acquire new knowledge for downstream tasks. The presence of the class gate leverages category information to aid the model in distinguishing class-specific knowledge during the learning process, thereby enhancing the effectiveness of the tailors.

As discussed earlier, learngenes serve as a superior backbone compared to pre-trained models due to their condensation of common knowledge. To explore this further, we analyze the properties of the common knowledge condensed in learngenes. Table 3 compares learngenes (w/o tailors) with pre-trained models on training tasks. The results reveal that learngenes demonstrate higher entropy, along with lower variance and kurtosis across different categories, indicating that the common knowledge they condense is largely class-agnostic. Such stability underscores that learngenes, as a backbone, provide superior adaptability to unfamiliar classes compared to pre-trained models.

We also visualize the learngenes with and without tailors in Figure 6. The visualizations show that the learngenes are not sensitive to category variations, consistently generating similar images across different class conditions. Although these images may lack detailed semantic information on their own, combining them with category-specific information (i.e., tailors) allows for the generation of images corresponding to specific categories, which further highlights the inherent commonality of knowledge within the learngenes.