FINE: Factorizing Knowledge for Initialization of Variable-sized Diffusion Models

TL;DR

FINE introduces a novel pre-training approach that factorizes knowledge into fundamental components called learngenes, enabling flexible initialization of variable-sized diffusion models without extensive retraining.

Contribution

FINE's method of representing model weights as a product of shared learngenes and layer-specific factors allows efficient, transferable initialization for models of different sizes, reducing training costs.

Findings

Achieves state-of-the-art performance in variable-sized model initialization.

Enables effective adaptation of models to diverse tasks.

Reduces retraining effort to light fine-tuning of layer-specific components.

Abstract

The training of diffusion models is computationally intensive, making effective pre-training essential. However, real-world deployments often demand models of variable sizes due to diverse memory and computational constraints, posing challenges when corresponding pre-trained versions are unavailable. To address this, we propose FINE, a novel pre-training method whose resulting model can flexibly factorize its knowledge into fundamental components, termed learngenes, enabling direct initialization of models of various sizes and eliminating the need for repeated pre-training. Rather than optimizing a conventional full-parameter model, FINE represents each layer's weights as the product of $U_{\star}$, $\Sigma_{\star}^{(l)}$, and $V_{\star}^\top$, where $U_{\star}$ and $V_{\star}$ serve as size-agnostic learngenes shared across layers, while $\Sigma_{\star}^{(l)}$ remains layer-specific. By jointly training these components, FINE forms a decomposable and transferable knowledge structure that allows efficient initialization through flexible recombination of learngenes, requiring only light retraining of $\Sigma_{\star}^{(l)}$ on limited data. Extensive experiments demonstrate the efficiency of FINE, achieving state-of-the-art performance in initializing variable-sized models across diverse resource-constrained deployments. Furthermore, models initialized by FINE effectively adapt to diverse tasks, showcasing the task-agnostic versatility of learngenes.