KOALA: Empirical Lessons Toward Memory-Efficient and Fast Diffusion Models for Text-to-Image Synthesis

• Large models incur substantial hardware costs, while unreleased training datasets make reproducing recent T2I results difficult for the open-source community.Source: §1 · p. 1
• Prior BK-SDM compression achieves a limited 33% compression rate on SDXL, and U-Net feature-distillation strategies had not been fully explored.Source: §1 · p. 2
• The proposed U-Nets reduce SDXL U-Net model size by up to 54% and 69%, while self-attention features are identified as crucial for distillation.The study also finds that high-resolution images with longer captions matter more than larger sets of low-resolution images with short captions, and that step-distilled teachers allow fewer denoising steps.Source: §1 · p. 2
• KOALA-Lightning-700M achieves 4× faster speed and 3× model efficiency than SDXL-Base while exhibiting satisfactory generation quality.KOALA-Lightning-700M also outperforms SSD-1B at 3× faster speed.Source: §1 · p. 2
• KOALA models operate on 8-GB GPUs, where SDXL models cannot be mounted.The authors present this as a practical advantage for practitioners in resource-constrained environments.Source: §1 · pp. 2–3

• Table 2 compares U-Net configurations; SDM-v2.0 uses 768² resolution, while SDXL and KOALA use 1024².Source: Table 2 · p. 3
• SDXL’s 2.56B U-Net is identified as the primary cause of its larger model size compared with SDM-v2.0’s 865M U-Net.Source: §2 · p. 3
• 83% of U-Net parameters are concentrated in ten transformer blocks at the lowest feature map, making them the main parameter bottleneck.Source: §2 · pp. 3–4
• SDXL distributes transformer blocks unevenly across stages, with more at lower feature levels and none at the highest levels described.Source: §2 · p. 3

3.1.1 Efficient U-Net architecture

• Layer-wise removal reduces transformer depth, addressing the limited compression of block removal alone in complex U-Nets.KOALA-700M reduces the lowest-level transformer layers from 10 to 5, while KOALA-1B reduces them to 6; KOALA-700M also removes the Mid block.Source: §3.1.1 · p. 4

3.1.2 Exploring Knowledge Distillation for SDXL

• Self-attention feature distillation gives the largest performance gain among tested feature types and helps distinguish objects, reducing appearance leakage.Decoder-stage self-attention distillation improves generation quality more than encoder-only self-attention distillation.Source: §3.1.2 · pp. 5–6

3.2 Lesson 2. Data: the impact of sample size, image resolution, and caption length

• Detailed captions improve text-image correspondence, while high-resolution images boost overall performance even with fewer training samples.The study selects LAION-POP, which combines high-resolution images with descriptive captions, as its main training data.Source: §3.2 · p. 6

3.3 Lesson 3. The influence of Teacher model

• Step-distilled SDXL-Turbo and SDXL-Lightning teachers yield KOALA-700M image quality comparable to or better than SDXL-Base with fewer denoising steps (10 vs. 25).The framework and Euler discrete noise schedule need no teacher-specific modifications.Source: Table 5 · p. 6, §3.3 · p. 7

4.1 Implementation details

• The models use LAION-POP and SDXL-Turbo or SDXL-Lightning as teachers, with 10-step FP16 inference using the Euler discrete scheduler.KOALA-Turbo and KOALA-Lightning use 512px and 1024px resolutions, respectively; training uses four NVIDIA A100 GPUs.Source: §4.1 · pp. 7–8

4.2 Main results

• KOALA-Lightning-700M/-1B outperform SDXL-Base-1.0 on HPSv2 and CompBench while running about 5×/4× faster, respectively.Against SDXL-Turbo and SDXL-Lightning, KOALA models have comparable or inferior HPSv2 but higher CompBench, with up to 3× smaller models and 1.7× lower memory usage.Source: §4.2 · p. 8, Table 6 · p. 8

4.3 Discussion

• With the same data setup, the proposed KD method achieves higher HPSv2 and CompBench scores than BK-SDM across different U-Net backbones.The authors attribute the result to distilling self-attention rather than only the last-layer feature, and report that this approach learns more discriminative representations for object and attribute distinctions.Source: §4.3 · pp. 8–9, Table 7 · p. 8
• Self-attention distillation beats other feature targets in DiT-M, while BK-style last-feature distillation performs worst.The study compresses DiT-XL from 28 to 14 layers while retaining the hidden dimension.Source: Table 8 · p. 8, §4.3 · p. 9
• KOALA backbones let PCM2,3 run faster with only a slight performance drop versus PCM-SDXL-Base, while generated images have comparable visual quality.The authors describe step-distillation and their KD approach as orthogonal; the qualitative comparison is in Fig. 20.Source: §4.3 · p. 9, Table 9 · p. 9

4.4 Model budget comparison on consumer-grade GPUs

• On 8GB GPUs, only KOALA models and SSD-Vega run among the compared models; KOALA-Lightning-700M is 2× faster than SSD-Vega.On 11GB GPUs, KOALA-Lightning-700M runs about 5× faster than SDXL-Base; the comparison uses the denoising steps in Table 6 and FP16.Source: Figure 7 · p. 9, §4.4 · p. 10

5 Limitations

• KOALA models have difficulty rendering legible text in generated images.The limitation is illustrated in Fig. 8 (Left).Source: §5 · p. 10
• KOALA models struggle with complex prompts involving multiple attributes or relationships between objects.The limitation is illustrated in Fig. 8 (Right).Source: §5 · p. 10
• The compression approach is heuristic and specific to SDXL’s heterogeneous U-Net, limiting its formulation as a general compression principle.The architecture combines convolutional and transformer blocks.Source: §5 · p. 10

6 Conclusion

• KOALA-Lightning-700M reduces SDXL-Base model size by 69% and latency by 79% while maintaining satisfactory generation quality.The authors propose KOALA models as cost-effective alternatives for practitioners with limited GPU resources.Source: §6 · p. 10
• The paper’s three empirical lessons concern self-attention in knowledge distillation, data characteristics, and teacher-model influence.The authors intend these lessons to benefit open-source efforts to improve T2I efficiency.Source: §6 · p. 10
• The authors plan to explore more general DiT compression methods and knowledge-distillation techniques based on self-attention distillation.They cite DiT’s architectural simplicity relative to SDXL’s U-Net as motivation for further exploration.Source: §6 · p. 10

A.1 Data

• After filtering blank text and corrupted images, LAION-Aesthetics V2 6+ contains 8,483,623 image-text pairs, but its prompts are notably brief.The authors note that brief prompts could hinder learning accurate image-text correspondence.Source: §A.1 · p. 15
• LLaVA-1.5 augments captions with image details, and the enriched captions are appended to original captions to retain contextual information such as proper nouns.The approach differs from replacing original captions with synthesized ones.Source: §A.1 · p. 15, Figure 9 · p. 15
• LAION-POP has 491,567 higher-resolution images and an average prompt length of 81; KOALA-Turbo and KOALA-Lightning are trained on it.Its images average 1274 × 1457 resolution, and its longer prompts are generated by CogVML and LLaVA-v1.5.Source: §A.1 · p. 15

A.2 Training

• The common protocol initializes the student U-Net from corresponding teacher blocks, freezes the text encoders, VAE, and teacher, and fine-tunes only the student U-Net.The main KOALA models are trained for 500K iterations on LAION-POP using four A100 GPUs and batch size 128.Source: §A.2 · p. 16
• Ablations use 30K iterations for distillation-strategy tests, 100K for data and teacher tests, and 50K for selected BK comparisons.The data tests use 1024 × 1024 images; the Turbo teacher test uses 512 × 512, while Base and Lightning teacher tests use 1024 × 1024.Source: §A.2 · p. 16

A.3 Inference

• KOALA-Turbo and KOALA-Lightning generate at 512 × 512 and 1024 × 1024, respectively, using Euler sampling, 10 denoising steps, and classifier-free guidance 3.5.Latency and memory comparisons omit torch.compile and quantization.Source: §A.3 · p. 17

A.4 Knowledge Distillation for Diffusion Transformer

• The DiT-M student reduces PixArt-Σ’s transformer depth from 28 to 14 layers and its size from 611M to approximately 313M parameters.It is initialized from corresponding DiT-XL layers and trained with task, output-KD, and feature-KD losses.Source: §A.4 · p. 17

A.5 Detailed formulation of training objectives

• The task loss trains reverse denoising, while output-level distillation matches student noise predictions to teacher predictions.The task loss uses sampled ground-truth noise and text conditioning; the teacher U-Net supplies the output-KD target.Source: §A.5 · p. 17

B.1 Self-attention based Feature-level Knowledge distillation

• SA-bottom performs best and SA-up worst among the tested self-attention locations, motivating distillation from consecutive early teacher blocks.The authors adopt SA-bottom in all experiments.Source: §B.1 · p. 18
• Adding last-feature guidance at SA-absent stages improves performance, with LF at DW-1 and UP-3 best; broader feature combinations do not improve it.The authors report that self-attention features are not complementary to the other tested features.Source: §B.1 · pp. 18–19

B.2 Attention visualization for Tab. 11a and Tab. 11b

• Self-attention maps progress from broad context to localized semantics, and at UP-1 the SA-based model attends to objects more distinctly than the LF-based model.The figure’s PCA visualizations compare rabbit and dog object attention; the caption also notes appearance leakage in BK-SDM outputs.Source: Figure 10 · p. 18, §B.2 · p. 19

B.3 Feature cosine similarity analysis for Tab. 11c

• Early transformer blocks contribute more substantially, while later blocks have less impact as features saturate across all analyzed stages.The analysis averages cross-layer cosine similarity over prompts from HPSv2; lower similarity indicates a greater contribution from the current block.Source: §B.3 · p. 19

B.4 Implementation details of SA-bottom

• SA-bottom pairs each student self-attention feature with consecutive teacher transformer blocks beginning at the first block.The illustrated teacher has depth 10 and the KOALA-1B student has depth 6.Source: §B.4 · pp. 19–20

C.1 Represenative prompts in Fig. 1

• Figure 1’s prompts include an astronaut raccoon, a koala making latte art, costumed animals, and elaborate character or scene descriptions.Examples include a witch-costumed cat and a richly specified cartoon hedgehog.Source: §C.1 · pp. 20–21

C.2 Comparison to other methods

• The qualitative comparison uses 1024 × 1024 images for KOALA-Lightning versus SDXL-Base, SDXL-Lightning, SSD-1B, and SSD-Vega, and 512 × 512 for KOALA-Turbo versus SDXL-Turbo.The cited figures identify the compared models and resolutions but do not state a comparative outcome.Source: §C.2 · p. 21, Figure 17 · p. 26, Figure 18 · p. 27

C.3 Comparison to BK-SDM

• Compared with BK-SDM, KOALA more consistently follows text prompts, while BK-SDM sometimes misses specified attributes or objects and produces structurally invalid images.The figure caption describes rabbit or dog appearance leakage in BK-SDM examples.Source: §C.3 · p. 21

D Failure cases

• KOALA-Lightning-1B struggles with legible text, human hands, and multi-attribute compositions; the authors conjecture LAION-POP may lack enough examples of these styles.The paper also notes difficulty with long text and unintended letters, and contrasts hand-rendering difficulty with stronger face representation.Source: D Failure cases · p. 21

E Societal Impacts

• KOALA applies an NSFW detector and discards generated images whose scores exceed a predefined threshold.The paper presents this filtering as a measure to address risks from harmful, violent, or adult content.Source: E Societal Impacts · pp. 21–22

2. Limitations

• KOALA’s reported limitations include generated-image failure cases and a simple, specific compression method.The passage does not specify the failure cases or further describe the method’s scope.Source: §2 · p. 31
• Training code and checkpoint weights are unavailable because of company policy, although the training datasets and evaluation benchmarks are downloadable.The authors say they plan to release trained weights after obtaining company permission.Source: §2 · p. 33
• Inference costs prevented repeated evaluations and error-bar reporting, despite evaluations using 2,400 HPSv2 images and 24,000 CompBench images.The authors state that they expect deviations to be minimal given the number of test images.Source: §2 · p. 33