Improvements to SDXL in NovelAI Diffusion V3 (2409.15997v2)

Abstract: In this technical report, we document the changes we made to SDXL in the process of training NovelAI Diffusion V3, our state of the art anime image generation model.

• The paper introduces v-prediction parameterization, enhancing numerical stability by dynamically balancing epsilon and x₀ predictions.
• The paper implements a zero terminal SNR noise schedule that significantly reduces artifacts and improves prompt compliance under high-noise conditions.
• The paper employs optimized training techniques, including higher noise levels and MinSNR loss weighting, to boost high-resolution image fidelity and overall consistency.

Improvements to SDXL in NovelAI Diffusion V3: An Expert Overview

The paper "Improvements to SDXL in NovelAI Diffusion V3" by Juan Ossa, Eren Doğan, Alex Birch, and F. Johnson presents a detailed technical report on the enhancements made to the SDXL model, which serves as the foundation for NovelAI's state-of-the-art anime image generation system. This essay aims to provide an insightful overview of the contributions and implications of this research.

Stable Diffusion eXtended (SDXL) has gained traction as an effective image generation model, benefitting from its open-source release. Leveraging this framework, the authors have introduced several key improvements to develop NovelAI Diffusion V3. The enhancements primarily focus on addressing parameterization, noise scheduling, and training optimizations.

v-Prediction Parameterization

The authors transitioned the SDXL model from an $\epsilon$-prediction to $v$-prediction parameterization. The motivation behind this change is multifaceted. Firstly, $\epsilon$-prediction becomes trivial at SNR=0, which limits its effectiveness. Conversely, $v$-prediction dynamically balances between $\epsilon$-prediction and $x_0$-prediction, facilitating robust performance across the entire noise range. This transition enhances numerical stability, minimizes color-shifting at higher resolutions, and accelerates convergence of sample quality.

Zero Terminal SNR

One of the standout contributions is the introduction of Zero Terminal SNR (ZTSNR). SDXL traditionally operates with a noise schedule that does not extend to pure-noise, leading to artifacts and non-prompt-relevant feature generation at inference time. By implementing a noise schedule that incorporates ZTSNR, the authors have trained the model to perform effectively even from states of infinite noise. As illustrated, this adjustment significantly reduces artifacts and improves the model's ability to generate prompt-compliant images, especially in high-noise regimes.

Improved High-Resolution Sampling

The introduction of higher maximum noise levels helps address the degradation of large features in high-resolution images. The authors suggest that doubling the maximum noise level is necessary to maintain signal-to-noise ratio (SNR) coherence as image resolution increases. Empirical results demonstrate that this method substantially enhances the fidelity of large-scale features.

MinSNR Loss Weighting

Addressing the multi-task nature of diffusion, the authors used MinSNR loss weighting to balance training across various timesteps. This approach enhances the model's capability to learn from both high-noise and low-noise conditions, thereby improving overall sample quality and consistency.

Dataset and Training

Utilizing a dataset of approximately 6 million images, predominantly anime-style illustrations, the model was trained on a 256x H100 GPU cluster for an extensive duration. The dataset was meticulously tagged, and sophisticated training practices were employed, including float32 precision with tf32 optimization. This rigorous training regime ensured thorough adaptation of the model to the dataset's specific characteristics.

Aspect-Ratio Bucketing

To address issues related to image cropping and aspect ratio inconsistencies, aspect-ratio bucketing was employed. This approach enables more natural and contextually appropriate image framing by grouping images of similar aspect ratios into buckets. This strategy ensures better token efficiency and reduces the probability of generating unnaturally cropped images.

Conditioning and Tag-based Loss Weighting

The V3 model continued using CLIP context concatenation for conditioning while introducing tag-based loss weighting to better handle rare concepts and mitigate the influence of overly common ones. This methodological enhancement ensures a balanced learning process, enabling the model to generate more diverse and accurate images.

VAE Decoder Finetuning

As in the previous version, the VAE decoder was fine-tuned to better produce anime textures, with an additional focus on eliminating undesirable JPEG artifacts. This finetuning helps in rendering more natural and high-quality anime-style images.

Empirical Results and Conclusions

The NovelAI Diffusion V3 produces coherent and relevant images at CFG scales between 3.5 and 5, which suggests superior data labeling and efficiency compared to standard SDXL settings. The model's capacity to generate a large volume of high-quality images daily highlights its practical efficacy.

In summary, NovelAI Diffusion V3 introduces significant technical advancements over the baseline SDXL architecture. The improvements in parameterization, noise scheduling, and training methodologies collectively enhance the model's image generation capabilities. Future work could explore further optimizing these techniques and applying them to other domains, thereby broadening the applicability and robustness of diffusion models in general.