PaperlayerView exact revision on arXiv

Source-linked AI summary

Not All Layers Are Created Equal: Adaptive LoRA Ranks for Personalized Image Generation

Donald Shenaj, Federico Errica, Antonio Carta

arXiv:2603.21884v1cs.CVcs.AIcs.LG

TL;DR

Fixed LoRA ranks require a costly layer-wise and subject-specific search, yet rank choice controls the balance between personalization quality and memory. LoRA2 learns ordered, component-specific ranks during fine-tuning. Across 29 subjects, it achieves a favorable DINO, CLIP-I, and CLIP-T trade-off with lower memory, including comparable scores to rank-512 models using 0.40 GB instead of up to 2.8 GB.

  • Problem

    Fixed LoRA ranks are chosen uniformly across components and subjects because full layer-wise, subject-specific rank selection has combinatorial cost, despite rank-dependent quality and memory trade-offs.

  • Method

    LoRA2 uses a differentiable variational framework that orders rank-index importance so each LoRA component can adapt its effective rank during fine-tuning.

  • Results

    Across 29 personalized subjects and two diffusion backbones, LoRA2 improves the trade-off among subject fidelity, text alignment, and memory; rank-512 models require up to 2.8 GB, versus 0.40 GB for comparable LoRA2 scores.

  • Takeaways & Limitations

    Adaptive ranks allocate capacity where it is beneficial, avoiding manual component-wise selection and the uniform-rank strategy empirically shown to be suboptimal.

  • Takeaways & Limitations

    Complex prompts can occasionally cause background colors to leak into the subject and subtly shift its appearance.

Abstract

from arXiv · show

Low Rank Adaptation (LoRA) is the de facto fine-tuning strategy to generate personalized images from pre-trained diffusion models. Choosing a good rank is extremely critical, since it trades off performance and memory consumption, but today the decision is often left to the community's consensus, regardless of the personalized subject's complexity. The reason is evident: the cost of selecting a good rank for each LoRA component is combinatorial, so we opt for practical shortcuts such as fixing the same rank for all components. In this paper, we take a first step to overcome this challenge. Inspired by variational methods that learn an adaptive width of neural networks, we let the ranks of each layer freely adapt during fine-tuning on a subject. We achieve it by imposing an ordering of importance on the rank's positions, effectively encouraging the creation of higher ranks when strictly needed. Qualitatively and quantitatively, our approach, LoRA$^2$, achieves a competitive trade-off between DINO, CLIP-I, and CLIP-T across 29 subjects while requiring much less memory and lower rank than high rank LoRA versions. Code: https://github.com/donaldssh/NotAllLayersAreCreatedEqual.

1 Introduction

Personalized diffusion models rely on LoRA, but fixed ranks across subjects and layers limit adaptation and waste memory. LoRA2 learns component-specific ranks during fine-tuning, improving the trade-off among image quality, text alignment, and memory.

  • Motivation: Personalized diffusion models fine-tune pretrained text-to-image models on a few sample images to generate specific subjects or styles, commonly using LoRA.LoRA injects low-rank updates into pretrained diffusion backbones for parameter-efficient personalization.
  • Motivation: A successful personalized model must balance subject or style quality, textual prompt fidelity, and low memory usage.These objectives are explicitly linked to the LoRA rank choice.
  • Problem: Fixed ranks across all components and subjects prevent layer-wise specialization, although optimal capacity varies with subject complexity, layer, and architecture.Complex subjects may need higher ranks, while simpler subjects and many layers can use substantially lower ranks.
  • Approach: LoRA2 adapts LoRA ranks during fine-tuning by ordering rank-index importance and encouraging each component toward its minimal effective rank.The approach is inspired by variational methods for learning adaptive neural-network width and avoids combinatorial manual search.
  • Results: Across 29 personalized subjects and two diffusion backbones, LoRA2 improves the trade-off among subject fidelity, text alignment, and memory consumption.Rank-512 models can require up to 2.8 GB of parameters, whereas LoRA2 achieves comparable scores with 0.40 GB.
  • Results: Optimal ranks vary significantly across subjects and layers, while regularizing rank parameters and LoRA weights yields compact models with minimal quality degradation.These findings support adaptive capacity allocation rather than a single globally fixed rank.

3 Method

LoRA² replaces globally fixed LoRA ranks with component-specific ranks learned during fine-tuning. A variational parameterization orders rank importance, enabling ranks to grow or shrink dynamically while targeting the minimum effective rank.

  • LoRA refresher: LoRA represents each update as B_ℓA_ℓ, where the low rank trades off adaptation flexibility against efficiency.The rank is typically fixed uniformly across model components.
  • Adaptive-rank variational framework: LoRA² introduces rank-index variables for each component, allowing additional weight sets to expand the LoRA matrices as rank increases.The variational framework treats rank as potentially infinite before selecting a finite effective rank.
  • Adaptive Rank LoRA: A learnable ν_ℓ determines the finite effective rank D_ℓ through the quantile of a discretized exponential distribution.The distribution assigns decreasing importance across rank indices, so later indices can be truncated with less effect.
  • Training loss: The resulting objective combines a mandatory fine-tuning loss with optional regularization over rank parameters and LoRA weights.As ν_ℓ changes, LoRA² recomputes component ranks and can introduce or remove parameters during training.
  • Implementation: LoRA² can grow or shrink dynamically during training, randomly initializing newly added B_ℓ and A_ℓ weights when the effective rank increases.Rescaled Kaiming initialization compensates for gradient effects caused by the rank-weighting matrix.
  • Implicit space search: Continuous optimization of ν replaces exhaustive search over S K^L subject-rank configurations with a single training process that softly adds or truncates ranks.This makes layer-wise and subject-specific architecture search more feasible despite optional regularization hyperparameters.

4 Experiments

The experiments evaluate LoRA² on personalized image generation using two diffusion backbones and 29 held-out subjects. They compare multiple fixed capacities and assess subject fidelity, text alignment, and generation efficiency.

  • Setup: Experiments use SDXL and KOALA-700m backbones with DreamBooth-based LoRA fine-tuning on 29 test subjects.One of 30 subjects was reserved for hyper-parameter tuning.
  • Setup: The study evaluates ranks from 8 through 512 for each subject and generates five images for each of ten prompts.LoRA² uses 500 SDXL training steps and 800 KOALA training steps, while fixed-rank LoRA uses 1000.
  • Evaluation: Performance is measured with DINO, CLIP-I, and CLIP-T scores comparing generated images with subject-image and prompt features.Scores are averaged first across generations and prompts for each subject, then across subjects.

5 Results

Across qualitative, aggregated, per-subject, rank-analysis, and ablation results, LoRA2 improves the trade-off among subject fidelity, text alignment, and memory usage. Its adaptive ranks preserve details while reducing unnecessary capacity.

  • 5.1 Qualitative Results: LoRA2 fully reproduces the clock and backpack subjects with high fidelity, preserving distinctive details that rank-512 LoRA misses.The clock’s numeral “3” and the backpack’s patch eye and tongue are preserved in LoRA2 outputs.
  • 5.2 Aggregated Results: Increasing fixed LoRA rank improves DINO and CLIP-I but lowers CLIP-T, revealing a trade-off between subject consistency and text alignment.Low-rank models omit distinctive attributes, whereas high-rank models better recognize subjects but increasingly deviate from textual descriptions.
  • 5.3 Per-Subject Performance: LoRA2 attains similar scores to rank 512 while using 0.40 GB instead of 2.80 GB, and it outperforms rank 64 on DINO and CLIP-I at comparable size.Compared with rank 64, LoRA2 has slightly lower CLIP-T but much higher DINO and CLIP-I scores.
  • 5.4 LoRA Rank Analysis: LoRA2 learns diverse ranks across layers and subjects: cross-attention more often reaches rank 512, self-attention tends toward lower ranks, and many components collapse to rank 1.This behavior allocates capacity differently across attention types while reducing ranks for components that need less adaptation.
  • 5.5 Ablation: 406 MB versus 2.7 GB: removing rank regularization greatly enlarges models, while removing entropy regularization lowers text alignment; both losses are needed for the intended trade-off.Without rank regularization, DINO and CLIP-I become marginally better, but file size increases substantially.

6 Conclusions

LoRA2 learns a suitable rank for each LoRA component, avoiding both combinatorial manual selection and globally fixed ranks. Across 29 subjects, it achieves a strong DINO, CLIP-I, and CLIP-T trade-off with lower memory consumption.

  • LoRA2 is an easy-to-implement, fully differentiable, model-agnostic modification of LoRA for personalized image generation.
  • Its ordered rank indices dynamically introduce or reduce each component’s rank according to the subject.
  • LoRA2 avoids manually selecting every component’s rank and empirically outperforms using one fixed rank for all components.
  • Across 29 subjects, LoRA2 achieves a strong trade-off among DINO, CLIP-I, and CLIP-T while using less memory.
  • The authors identify multi-subject, model-merging, and larger-diffusion-model settings as future directions.

A1 Additional Implementation Details

The implementation trains 1024 × 1024 models with batch size 1 and mixed-precision, memory-saving optimization techniques on 64GB NVIDIA Ampere A100 GPUs.

  • Models were trained at 1024 × 1024 resolution with batch size 1 and learning rate 5 × 10^-5.
  • Training used fp16 mixed precision, gradient checkpointing, and 8-bit Adam optimization.
  • Experiments ran on NVIDIA Ampere A100 GPUs with 64GB RAM.

A2 Prompts

Appendix Table A1 lists the full evaluation prompts, pairing personalized subjects with varied contexts, attributes, and environments.

  • Table A1 provides the full prompts used for evaluation.
  • Evaluation contexts include snow, beaches, forests, city parks, outer space, studios, shelves, and reflective surfaces.
  • Prompt attributes include wet, shiny, glowing, vintage, cube-shaped, modern minimalistic, and rainbow-scarf variants.
  • Several prompts place subjects in distinctive scenarios such as chef outfits, police uniforms, wizard robes, sunglasses, hats, and monocles.

A3 Full Self-Attention and Cross-Attention Ranks

The appendix reports full self-attention and cross-attention rank visualizations for five distinct subjects under SDXL and KOALA-700m.

  • SDXL self-attention ranks are shown for five distinct subjects.
  • SDXL cross-attention ranks are shown for five distinct subjects.
  • KOALA-700m self-attention ranks are shown for five distinct subjects.
  • KOALA-700m cross-attention ranks are shown for five distinct subjects.

A4 KOALA Per-Class Scores

KOALA-700m results show that the best LoRA rank varies by subject, with substantial variability in per-subject rank selection.

  • The optimal LoRA rank changes depending on the subject in KOALA-700m experiments.Per-subject scores are reported with rank 64 marked as the commonly used default and each subject’s best value highlighted.

A5 Additional Qualitative Results

Additional SDXL results show LoRA2 preserving fine-grained subject details and handling complex prompts more effectively than fixed-rank LoRA.

  • LoRA2 consistently reproduces the can label across generated images, indicating stronger fidelity to fine-grained subject details.
  • LoRA2 generalizes to complex prompt scenarios, while fixed-rank LoRA often fails to recontextualize subjects properly.

A6 Limitations

The paper identifies scope and deployment limitations for LoRA2, including untested style learning, adapter-rank incompatibility during merging, and occasional appearance artifacts under complex prompts.

  • The evaluation focuses on personalized subject learning, leaving style learning as future work.
  • Different adapter ranks across subjects complicate model merging because lower-rank adapters must be expanded before merging.Composition-based approaches can avoid explicit adapter merging.
  • Complex prompts can cause background colors to leak into the subject and subtly shift its appearance.The artifact also appears in competing approaches, although LoRA2 retains superior subject fidelity under challenging prompts.
Loading 2603.21884v1…