PaperlayerView exact revision on arXiv

Source-linked AI summary

SANA 1.5: Efficient Scaling of Training-Time and Inference-Time Compute in Linear Diffusion Transformer

Enze Xie, Junsong Chen, Yuyang Zhao, Jincheng Yu, Ligeng Zhu, Chengyue Wu, Yujun Lin, Zhekai Zhang, Muyang Li, Junyu Chen, Han Cai, Bingchen Liu, Daquan Zhou, Song Han

arXiv:2501.18427v4cs.CV

TL;DR

Scaling text-to-image diffusion models improves quality but creates substantial training and inference costs. SANA-1.5 combines knowledge-preserving model growth, depth pruning, memory-efficient optimization, and repeated-sample inference scaling to address these costs. It reports 0.81 on GenEval, improved to 0.96 through inference scaling, while supporting flexible model deployment under constrained resources.

  • Problem

    Text-to-image generation has increasingly relied on larger models, creating substantial training and inference costs, while the scalability of efficient linear diffusion transformers remains an open question.

  • Method

    SANA-1.5 combines pretrained-model depth growth, block-importance-based depth pruning, repeated sampling with VLM selection, and the memory-efficient CAME-8bit optimizer.

  • Results

    0.81 GenEval score is achieved through training scaling and improves to 0.96 with inference scaling, while training convergence is 2.5× faster than scaling up and training from scratch.

  • Takeaways & Limitations

    The framework provides multiple routes to high-quality text-to-image generation under different compute budgets, making large-scale model training and deployment more accessible.

  • Takeaways & Limitations

    Challenges remain in complex generation tasks, particularly text rendering and human details.

Abstract

from arXiv · show

This paper presents SANA-1.5, a linear Diffusion Transformer for efficient scaling in text-to-image generation. Building upon SANA-1.0, we introduce three key innovations: (1) Efficient Training Scaling: A depth-growth paradigm that enables scaling from 1.6B to 4.8B parameters with significantly reduced computational resources, combined with a memory-efficient 8-bit optimizer. (2) Model Depth Pruning: A block importance analysis technique for efficient model compression to arbitrary sizes with minimal quality loss. (3) Inference-time Scaling: A repeated sampling strategy that trades computation for model capacity, enabling smaller models to match larger model quality at inference time. Through these strategies, SANA-1.5 achieves a text-image alignment score of 0.81 on GenEval, which can be further improved to 0.96 through inference scaling with VILA-Judge, establishing a new SoTA on GenEval benchmark. These innovations enable efficient model scaling across different compute budgets while maintaining high quality, making high-quality image generation more accessible. Our code and pre-trained models are released.

1. Introduction

SANA-1.5 addresses the high computational cost of scaling text-to-image diffusion models through coordinated training-time growth, depth pruning, and inference-time scaling. These strategies improve efficiency and quality while supporting deployment across different resource budgets.

  • Motivation: Text-to-image models have grown substantially, but larger model sizes impose prohibitive training and inference costs.Recent industry models span 0.6B to 24B parameters.
  • Model Growth: SANA-1.5 scales SANA from 1.6B to 4.8B parameters by initializing a large model from a pretrained small model and strategically adding blocks.The approach reuses knowledge from the smaller model instead of training the larger model from scratch.
  • Model Pruning: Model depth pruning analyzes block importance and fine-tunes compressed 60-block models into 40-, 30-, or 20-block configurations.Fine-tuning can recover quality in as little as 5 minutes on a single GPU.
  • Inference Scaling: Inference-time scaling generates multiple samples and uses a VLM-based verifier to select the best candidate, allowing smaller models to match larger-model quality through additional computation.This strategy trades inference computation for model capacity.
  • Implications: The combined growth, pruning, and inference-scaling framework supports comparable or better results with significantly reduced training resources.The paper frames optimization trajectories and resource-aware deployment as alternatives to simply increasing parameter count.
  • Results: 0.81 GenEval score is achieved through training scaling, rising to 0.96 with inference scaling.The paper reports 2.5× faster training convergence than scaling up and training from scratch, and presents the 0.96 result as a new state-of-the-art on GenEval.

2. Methods

SANA-1.5 scales efficiently through model growth, memory-efficient optimization, depth pruning, and inference-time sampling. These methods reuse learned knowledge, compress models, and use additional sampling candidates with VLM-based selection.

  • Efficient Model Growth: The partial preservation strategy keeps the first N pretrained layers and randomly initializes the additional M layers using a normal distribution.The paper compares partial preservation, cyclic replication, and block replication initialization.
  • Memory-Efficient CAME-8bit Optimizer: CAME-8bit combines matrix-factorized second-order moments with block-wise 8-bit quantization of first-order moments.Only large matrices in linear and 1×1 convolution layers are quantized, while critical statistics remain in 32-bit precision.
  • Memory-Efficient CAME-8bit Optimizer: The optimizer reduces memory usage to approximately 1/8 of AdamW while preserving CAME’s convergence properties.Actual savings are lower than the approximation because small layers, second-order statistics, and quantization metadata remain costly.
  • Model Depth Pruning: Depth pruning analyzes input-output similarity across diffusion blocks and removes less important blocks to create smaller model configurations.Importance is averaged across diffusion timesteps and a calibration dataset of 100 diverse prompts; quality can be quickly recovered through fine-tuning.
  • Inference-Time Scaling: Inference-time scaling generates multiple samples and uses a VLM-based tournament to select images that best match the prompt.The approach treats additional sampling computation as a complement to parameter scaling, while denoising-step scaling offers only minor improvements.

3. Experiments

Experiments evaluate SANA-1.5 across training, pruning, optimization, and inference scaling. The results show strong quality, efficiency, and compression across model sizes and compute settings.

  • Experimental setup: SANA-4.8B scales to 60 layers while retaining SANA-1.6B’s channel and FFN dimensions, with training conducted using distributed PyTorch DDP.The model uses 64 NVIDIA A100 GPUs across 8 DGX nodes.
  • Inference scaling: SANA-1.5 + inference scaling with 2048 samples achieves absolute SoTA on the GenEval benchmark against open-source and commercial methods.The benchmark uses Playground v3 values for baseline methods.
  • Main results: 0.81 GenEval is achieved after post-training, outperforming Playground v3’s 0.76 while maintaining lower latency and higher throughput than FLUX-dev.The speed comparison is measured on one A100 GPU with FP16 precision.
  • Model pruning: 0.672 versus 0.664 GenEval shows that a pruned and fine-tuned model outperforms a model trained from scratch at small computational cost.Table 3 evaluates directly pruned 3.2B and 1.6B models, with “+FT” denoting efficient fine-tuning.
  • Inference scaling: 0.96 overall GenEval accuracy is reached by selecting samples from 2048 generated images, up from 0.81 for naive single-image generation.Position improves from 0.59 to 0.96, while Color Attribution improves from 0.65 to 0.87.
  • Optimization: CAME-8bit reduces memory consumption by 25% compared to AdamW while preserving convergence speed, using 43GB versus 57GB.The optimizer comparison is conducted for SANA-1.6B training.
  • Model pruning: Adaptive depth pruning preserves semantic content across 1.6B–4.8B configurations, while brief fine-tuning recovers overall image quality after aggressive pruning.The caption reports recovery after 100 steps on one GPU.
  • Analysis: Block importance analysis identifies lower importance in middle-to-end blocks of the model grown to 4.8B, informing pruning and initialization choices.The analysis compares SANA-1.6B, a 4.8B model trained from scratch, and the final model with initialization.

4. Related Work

Related work spans rapidly improving diffusion architectures, efficiency-focused text-to-image models, autoregressive alternatives, and scaling-law research. Prior studies also investigate preference alignment and quality improvements without increasing model size.

  • Text-to-image architectures: Text-to-image generation has progressed from Stable Diffusion to DiT and related architectures.Efficiency-focused examples include SnapGen, PixArt-α, and SANA.
  • Alternative architectures: Autoregressive models have developed rapidly and achieve quality comparable to diffusion models.
  • Scaling and alignment: Research in language and vision has revealed power-law relationships, while image-generation studies have explored RLHF for human-preference alignment.
  • Efficient improvement: Concurrent work explores improving generation quality without increasing model size.

5. Conclusion

SANA-1.5 addresses training and inference efficiency through complementary scaling strategies. Together, these methods improve quality under limited computing budgets and broaden access to large-scale generative-model research.

  • Contributions: CAME-8bit and stable model growth address training efficiency, while repeat sampling and depth pruning address inference scaling and acceleration.
  • Conclusion: The combined approaches enable significant quality improvements under limited computing budgets.
  • Conclusion: The paper positions these methods as making large-scale generative models more accessible to researchers with limited resources.

Impact Statement

SANA-1.5 addresses misuse risks with a safety check model while acknowledging remaining challenges in complex image generation. The paper highlights text rendering and human details as unresolved areas.

  • Safety: A safety check model screens user prompts for NSFW content before generation.The prompt is first sent to the safety check model to determine whether it contains NSFW content.
  • Open Challenges: Complex generation tasks remain challenging, particularly text rendering and human details.
  • Responsible Deployment: The work encourages responsible development and deployment to prevent misuse.

A. Full Related Work

Related work traces text-to-image progress from diffusion models and DiTs toward more efficient training, pruning, and inference scaling. Prior studies establish scaling patterns across model size, data, compute, and sample count.

  • Text-to-Image Generation: Text-to-image generation evolved from Stable Diffusion toward Diffusion Transformers, alongside efforts to reduce training costs.PixArt-α is cited as achieving competitive quality with training requirements equal to 10.8% of Stable Diffusion v1.5.
  • Diffusion Model Pruning: Diffusion-model pruning methods reduce model size and computation for resource-constrained deployment.Prior approaches include neural architecture search, elastic depth, and distillation, while this work uses a heuristic to identify less important DiT blocks.
  • Training Scaling: Training-scaling research reports power-law relationships between model accuracy and model size, dataset size, and compute.These scaling patterns have been observed across several orders of magnitude in language and vision domains.
  • Inference Scaling: Inference-scaling studies find that solution coverage can grow log-linearly with the number of samples.Self-consistency methods further improve accuracy by sampling multiple reasoning paths and selecting the most consistent answer.

B. Inference-Time Scaling Details

The inference-time scaling setup fine-tunes NVILA on generated image-prompt pairs and evaluates alignment using GenEval. The supporting setup also documents initialization strategies and training hyperparameters.

  • Dataset: 2.5M generated images are paired with 15,654 unique prompts and evaluated for alignment using GenEval.The prompts cover 80 object classes with attributes including color, spatial location, quantity, and relationships.
  • Initialization: Figure 11 compares Partial Preservation, Cyclic Replication, and Block Replication initialization strategies.Partial Preservation keeps pretrained layers and randomly initializes new layers; the other strategies replicate pretrained layers in different patterns.
  • Training Setup: The training setup uses learning rate 2 × 10^-5, Adam, cosine scheduling with warmup ratio 0.03, and batch size 8 per device.The model is trained on the 2M dataset for one epoch.

C. More Implementation Details

Implementation details cover stabilized attention, multilingual captioning, initialization choices, and comparisons of model growth with training from scratch. Together, these details describe the engineering choices supporting efficient scaling.

  • Attention Design: RMSNorm is applied to queries and keys in both linear self-attention and vanilla cross-attention.The design targets training stability in large diffusion models by addressing growing attention logits and FP16 numerical instability.
  • Multilingual Data: A multilingual auto-labeling pipeline creates Chinese, English-Chinese mixed, and emoji-enriched captions from 100K English prompts.GPT-4 translates the prompts to build training data with diverse descriptions of the same visual content.
  • Model Growth: Model growth outperforms training from scratch across FID, CLIP score, and DPGBench metrics.It achieves better results within the same duration or equivalent performance with approximately 60% less training time.
  • Initialization: Partial Preservation initialization preserves pretrained layers and randomly initializes new layers.The paper adopts it for simplicity and stability, reporting more stable training dynamics than cyclic and block expansion strategies.

D. More Results

SANA-1.5 improves scaling efficiency, preserves quality across pruned model sizes, and demonstrates broad generation capabilities. Its inference-time selection and prompt rewriting further enhance prompt matching, detail, and visualization quality.

  • Model Growth Results: 60% less training time achieves equivalent quality versus training from scratch across FID, CLIP score, and DPG benchmarks.The model-growth strategy also achieves better quality within the same training duration.
  • Pruned Model Sizes: 0.693 GenEval for SANA-1.5 4.8B remains 0.684 at 3.2B and 0.672 at 1.6B, exceeding SANA-1.0 1.6B at 0.665.Pruning supports flexible model-size selection while preserving strong accuracy.
  • Visualization Results: SANA demonstrates high-fidelity detail rendering, semantic understanding, and reliable text generation across diverse scenarios.Examples include intricate textures, complex compositions, accurate text rendering, and faithful prompt interpretation.
  • Inference-Time Scaling: VLM judges filter prompt-mismatching samples, including images with incorrect object counts, during inference-time scaling tournaments.When similarly high-quality images both match the prompt, the VLM can judge that both images match the prompt.
  • Prompt Rewrite Enhancement: GPT-4o prompt rewriting adds details that produce richer and more detailed visualizations at inference time.The comparison uses a simple prompt about two sheep and an expanded pastoral description.

E. Discussion of Potential Misuse of SANA-1.5

SANA-1.5 is paired with safety checking intended to address misuse involving NSFW prompts. The section also presents multilingual, scaling, high-resolution, and prompt-enhancement examples illustrating the system’s broader capabilities.

  • Safety Check Pipeline: NSFW prompts are routed to a safety check model before generation, and requests classified as NSFW are rejected.Non-NSFW prompts continue to SANA-1.5 for image generation.
  • Multilingual Capabilities: Fine-tuning with 100k multilingual samples supports mixed English, Chinese, and emoji prompts compared with English-only training.The comparison is presented in Figure 15.
  • Safety Check Pipeline: ShieldGemma reportedly filters user-entered NSFW prompts under strict thresholds after extensive testing.The passage states that the pipeline will not create harmful AI-generated content.
  • Model Scaling: SANA-1.5 models of different sizes show scaling effects, with larger models enhancing detail rendering, image quality, and semantic alignment.Figure 16 compares SANA-1.5 and SANA-1.0 across model sizes.
  • Inference-Time Scaling: Inference-time scaling uses VLM judges to filter prompt-mismatching images during sample selection.Figure 17 visualizes this tournament-based process.
  • High-Resolution Generation: High-resolution examples cover accurate prompt following, spatial reasoning, text rendering, and aesthetics across styles and aspect ratios.These examples are shown in Figure 18.
  • Prompt Enhancement: GPT-4o prompt enhancement produces richer visual quality and details than outputs from original simple prompts.Figure 19 compares generation before and after prompt rewriting.
Loading 2501.18427v4…