ethanfel-ComfyUI-LoRA-Optimizer/docs/acestep-merge-optimization.md

ACE-Step 1.5 LoRA Merge Optimization

Technical reference for all architecture-specific choices made in ZImage LoRA Merger for ACE-Step 1.5 music generation model LoRAs.

ACE-Step 1.5 Architecture Overview

ACE-Step 1.5 is a 2B-parameter Diffusion Transformer (DiT) for music generation. It uses a Qwen3-based architecture with 24 transformer blocks, each containing self-attention, cross-attention, and MLP (SwiGLU) sub-modules.

Model Configuration

Parameter	Value
hidden_size	2048
num_hidden_layers	24
num_attention_heads	16
num_key_value_heads	8 (GQA)
head_dim	128
intermediate_size	6144
sliding_window	128
text_hidden_dim	1024
timbre_hidden_dim	64
dtype	bfloat16

Source: ACE-Step/Ace-Step1.5 config.json

Hybrid Attention Pattern

Self-attention alternates between two modes across the 24 layers:

• Odd layers (0, 2, 4, ..., 22): Sliding Window Attention (window=128). Captures local acoustic nuances — timbre texture, transients, articulation.
• Even layers (1, 3, 5, ..., 23): Global Group Query Attention (GQA). Maintains long-range rhythmic and melodic consistency.

Cross-attention always uses full/global attention regardless of layer index. This ensures voice identity from the timbre encoder has unrestricted access to every position in the audio sequence.

Source: ACE-Step 1.5 paper (arXiv:2602.00744), Architecture analysis (zenn.dev)

Conditioning Stack

Three conditioning pathways are concatenated along the token axis and injected via cross-attention into every DiT layer:

1. Qwen3-0.6B caption embeddings — text descriptions of genre, mood, style
2. AceStepLyricEncoder (8-layer bidirectional transformer) — lyric content
3. AceStepTimbreEncoder (4-layer bidirectional transformer) — voice/instrument timbre from reference audio (30s context window, 25Hz latent via 1D VAE)

The timbre encoder processes reference audio into a 2048-dimensional embedding. When no reference audio is provided, 30 seconds of silence is used as input. Non-padded tokens from all three sources are packed left before concatenation to prevent padding artifacts in attention.

Source: ACE-Step 1.5 model code, ACE-Step 1.5 Explained (Substack)

DiT Layer Structure (AceStepDiTLayer)

Each of the 24 layers contains:

AceStepDiTLayer:
  self_attn_norm  (Qwen3RMSNorm)
  self_attn       (AceStepAttention — sliding or global per layer_types)
  cross_attn_norm (Qwen3RMSNorm)
  cross_attn      (AceStepAttention — always global)
  mlp_norm        (Qwen3RMSNorm)
  mlp             (Qwen3MLP — SwiGLU: gate_proj, up_proj, down_proj)
  scale_shift_table (nn.Parameter, shape [1, 6, hidden_size] — AdaLN modulation from timestep)

The scale_shift_table provides 6 modulation values per layer: (shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa) used for adaptive layer norm on self-attention and MLP outputs.

Source: ACE-Step 1.5 model code (xl_turbo)

LoRA Key Structure (v1.5)

v1.5 uses PEFT-format LoRA adapters. Canonical key patterns after normalization:

Module	Key Pattern	Shape (lora_A)	Shape (lora_B)
Self-attention	diffusion_model.layers.{0-23}.self_attn.{q,k,v,o}_proj	[r, 2048]	q/o: [2048, r], k/v: [1024, r]
Cross-attention	diffusion_model.layers.{0-23}.cross_attn.{q,k,v,o}_proj	[r, 2048]	q/o: [2048, r], k/v: [1024, r]
MLP	diffusion_model.layers.{0-23}.mlp.{gate,up,down}_proj	gate/up: [r, 2048], down: [r, 6144]	gate/up: [6144, r], down: [2048, r]

Note: k/v projections output 1024 (8 KV heads x 128 head_dim) due to GQA, while q/o projections output 2048 (16 heads x 128 head_dim).

Community LoRAs commonly target all three module types (self_attn, cross_attn, mlp) and may use DoRA with per-module rank patterns via Side-Step's Fisher Information analysis.

Source: HuggingFace community LoRAs, Side-Step toolkit

Voice-Critical Layer Analysis

When merging a vocal LoRA with a music/style LoRA, voice identity is at risk of being diluted. Understanding which layers carry voice information guides the merge strategy.

Priority ranking for voice preservation

Priority	Module type	Reasoning
Highest	cross_attn (all 24 layers)	Primary voice-identity pathway. Attends to timbre encoder output (voice characteristics) concatenated with text/lyric conditioning. Always uses full global attention — never sliding window.
High	self_attn in sliding window layers (odd: 0, 2, 4, ...)	Captures local acoustic details — timbre texture, vocal transients, articulation within 128-token windows.
Medium	self_attn in global layers (even: 1, 3, 5, ...)	Long-range structure — melody contour, rhythmic patterns, song form. More style than voice.
Lower	mlp (all layers)	General audio feature transformation. Less specific to voice identity.

Why cross-attention is critical

In ACE-Step 1.5, the model operates on a "unified conditioning stack where Qwen3-0.6B caption embeddings are concatenated with dedicated timbre and lyric encoders and injected via Cross-Attention" (arXiv:2602.00744v3, Section 3.2). Cross-attention is the primary injection point for voice conditioning, though once injected, voice information propagates through self-attention and MLP layers as well.

When a vocal LoRA is trained, its cross-attention weights learn to respond strongly to the voice's timbre representation. When a music LoRA is trained on instrumental data, its cross-attention weights learn to respond to genre/style conditioning. These are typically orthogonal — they encode different aspects of the conditioning signal.

Naive averaging or TIES trimming on cross-attention layers destroys the vocal LoRA's learned response to timbre conditioning.

Source: ACE-Step 1.5 paper (arXiv:2602.00744v3), Voice cloning discussion (Issue #259), Architecture analysis (zenn.dev)

Merge Optimization Decisions

1. Dedicated acestep_dit Architecture Preset

Decision: Create a separate preset instead of sharing the generic dit preset.

Justification: The generic dit preset is tuned for image/video diffusion transformers (Flux, WAN, Z-Image, LTX) where all attention layers serve similar functions. In ACE-Step, cross-attention has a fundamentally different role (voice conditioning) that requires more conservative merge behavior.

Preset values and their rationale:

ties_conflict_threshold: 0.25 → 0.35

The TIES conflict threshold determines when the merger switches from weighted_average/SLERP to TIES (trim-elect-merge) mode. TIES aggressively trims low-magnitude values and resolves sign conflicts by majority vote.

For vocal + music LoRA merging, TIES trimming is destructive to voice:

• Cross-attention layers encode orthogonal concepts (voice timbre vs genre)
• TIES trims the minority-magnitude signal, which could be the voice
• Sign conflicts in cross-attn are expected (different conditioning responses) but don't represent real semantic conflict

Raising the threshold from 0.25 to 0.35 means TIES only activates for genuinely high conflict (>35% excess), causing more prefixes to use SLERP (magnitude-preserving interpolation) instead.

orthogonal_cos_sim_max: 0.25 → 0.30

This threshold determines when two LoRAs are classified as "orthogonal" (independent signals). Orthogonal LoRAs get routed to weighted_average → SLERP instead of TIES, since ~50% sign conflict is the statistical baseline for independent vectors and doesn't represent real conflict.

Vocal and music LoRAs are typically orthogonal in most layers (cosine similarity near zero). Widening the orthogonal band from 0.25 to 0.30 catches more of these cases, ensuring SLERP is used to preserve both voice and music signals at full magnitude.

orthogonal_conflict_max: 0.60 → 0.65

The maximum conflict ratio allowed for the orthogonal classification. Slightly raised to accommodate the naturally higher base-rate conflict seen in audio LoRA pairs.

dare_ideal_density: 0.8 → 0.85

When DARE sparsification is used, this controls the target density (fraction of parameters kept). Higher density preserves more signal, which is important for voice quality — dropping 15% of voice-encoding parameters is less destructive than dropping 20%.

auto_strength_orthogonal_floor: 0.85 → 1.0

When vocal and music LoRAs are orthogonal, auto-strength normally attenuates the merge to prevent magnitude explosion. For ACE-Step, full magnitude preservation (floor = 1.0) is critical because:

• Audio LoRAs encode energy/volume information in weight magnitude
• Attenuation reduces vocal presence in the generated output
• Orthogonal signals don't interfere, so no magnitude reduction is needed

This matches the behavior already applied to video architectures (WAN, LTX) via _VIDEO_ARCH_ORTHOGONAL_FLOOR, where motion energy similarly requires magnitude preservation.

Source for video precedent: existing _VIDEO_ARCH_ORTHOGONAL_FLOOR in codebase

2. Orthogonal Floor Override

Decision: Add ACE-Step to _VIDEO_ARCH_ORTHOGONAL_FLOOR with value 1.0.

Justification: Same reasoning as the preset floor. This is a secondary enforcement point that catches edge cases where the preset floor might be overridden by other logic. Ensures orthogonal vocal + music LoRA merges always preserve full magnitude.

3. v1.0 Format Detection and Normalization

Decision: Support ACE-Step v1.0 LoRA format alongside v1.5.

Justification: v1.0 LoRAs use a completely different key structure (transformer_blocks.N.attn.to_q vs v1.5's layers.N.self_attn.q_proj) and include unique voice-critical keys (speaker_embedder, lyric_encoder) that are absent in v1.5.

Without v1.0 support, these LoRAs were misdetected as WAN (because transformer_blocks.0.cross_attn matches the WAN detector's blocks. check) and would be incorrectly normalized.

v1.0 → v1.5 Key Mapping

v1.0 pattern	v1.5 canonical form
transformer_blocks.N	layers.N
.attn. (bare)	.self_attn.
.to_q. / .to_k. / .to_v.	.q_proj. / .k_proj. / .v_proj.
.to_out.0.	.o_proj.
speaker_embedder	diffusion_model.speaker_embedder
lyric_encoder.encoders.N.self_attn.linear_q	diffusion_model.lyric_encoder.encoders.N.self_attn.q_proj

Source: ACE-Step v1.0 repo, Chinese RAP LoRA

v1.0 Voice-Critical Keys

• speaker_embedder: Directly encodes voice timbre as a 512-dim embedding projected to 2560-dim for cross-attention. Trained with 50% dropout in the base model, so the remaining signal is concentrated and critical. During merge, if only the vocal LoRA targets this prefix, it becomes a single-LoRA case → weighted_sum → full preservation.

• lyric_encoder: 6-block conformer encoding lyric content. Less voice-specific but encodes pronunciation/phonetic patterns that contribute to vocal character.

Source: ACE-Step v1.0 paper (arXiv:2506.00045)

4. Bug Fix: to_out → o_proj Mapping

Decision: Fix the regex that mapped to_out → out_proj (wrong) to correctly produce o_proj, and handle v1.0's to_out.0. suffix.

Before (broken):

re.sub(r"\.to_(q|k|v|out)\.", lambda m: f".{m.group(1)}_proj.", new_k)
# to_out → out_proj (WRONG — v1.5 uses o_proj)
# to_out.0 → not matched (the .0. after to_out breaks the pattern)

After (fixed):

re.sub(r"\.to_(q|k|v)\.", lambda m: f".{m.group(1)}_proj.", new_k)
re.sub(r"\.to_out\.0\.", ".o_proj.", new_k)
re.sub(r"\.to_out\.", ".o_proj.", new_k)

This bug would cause v1.0 LoRA output projection keys to not match v1.5 keys, breaking cross-LoRA prefix grouping when merging a v1.0 LoRA with a v1.5 LoRA.

5. WAN Detection Refinement

Decision: Exclude transformer_blocks from the WAN architecture detector.

Justification: WAN uses bare blocks.N while ACE-Step v1.0 uses transformer_blocks.N. The substring blocks. appears in both, causing v1.0 ACE-Step LoRAs to be misdetected as WAN. The fix checks for transformer_blocks and skips the WAN path.

How the Existing Merge System Benefits ACE-Step

The per-prefix adaptive merge strategy already provides significant advantages over the community's standard merge tool (MERGE-LORA.py), which performs blind weighted-sum across all layers.

Per-prefix strategy selection

For each LoRA key prefix (e.g., layers.5.cross_attn.q_proj), the merger:

1. Computes pairwise conflict metrics (sign conflict, cosine similarity, excess conflict, subspace overlap, magnitude ratio)
2. Classifies the relationship (orthogonal, aligned, conflicting)
3. Selects the optimal strategy:
   • SLERP for orthogonal signals (preserves both at full magnitude)
   • Consensus for aligned signals (Fisher-proxy importance weighting)
   • TIES for genuine conflicts (trim + elect sign + disjoint merge)
   • Weighted sum for single-LoRA prefixes

Expected behavior for vocal + music LoRA merge

Prefix type	Expected relationship	Expected strategy	Voice outcome
cross_attn	Orthogonal (voice vs genre conditioning)	SLERP	Both voice and genre preserved
self_attn (odd/sliding)	Low conflict (different local patterns)	SLERP or weighted_average	Timbre details preserved
self_attn (even/global)	Variable	Auto-selected	Structure blended
mlp	Low-medium conflict	SLERP or weighted_average	General features blended
speaker_embedder (v1.0)	Single-LoRA (only vocal targets it)	Weighted sum	Full voice preservation

Comparison with community MERGE-LORA.py

Feature	MERGE-LORA.py	ZImage LoRA Merger
Merge method	Weighted sum only	TIES, SLERP, consensus, weighted avg, auto-select
Conflict handling	None (opposing signals cancel)	Per-prefix conflict analysis
Architecture awareness	None	ACE-Step detected, acestep_dit preset applied
Per-layer control	None (same weights everywhere)	Per-prefix adaptive strategy
Voice preservation	Manual strength tuning only	Automatic via orthogonal detection + SLERP
Magnitude preservation	Optional RMS clamp	Auto-strength with orthogonal floor = 1.0
Rank handling	SVD truncation to min(r1, r2)	Configurable SVD compression
Multi-LoRA	2 at a time (chain for 3+)	N-way native

Source for MERGE-LORA.py: DisturbingTheField HuggingFace repos

Practical Recommendations for ACE-Step Users

Merging ACE-Step LoRAs

1. Load LoRAs into a LoRA Stack node and connect to a LoRA Optimizer node (or use LoRA AutoTuner for automatic quality optimization)
2. Architecture preset will auto-detect as acestep_dit
3. Use per_prefix optimization mode (default)
4. Keep normalize_keys enabled (handles v1.0/v1.5/Kohya/PEFT format differences)
5. Set LoRA scale to 0.3–0.5 at inference for best results

This applies to all ACE-Step LoRA merges — vocal + music, genre + genre, instrument + style, etc. The acestep_dit preset is tuned for preserving distinct signals (voice, genre, instrument character) that would otherwise be diluted by naive averaging.

If voice quality degrades (vocal + music merges)

• Try merge_strategy_override: slerp to force SLERP on all prefixes
• Reduce the music LoRA's strength relative to the vocal LoRA
• Check the merge report — cross_attn prefixes should show SLERP, not TIES

Training tips for better merge compatibility

These are from the ACE-Step community, not from our merge code:

• Don't isolate vocals before training — keep instrumentals in training data
• Use trigger words to tag the singer consistently
• Reduce ssl_coeff (mHuBERT loss weight) to 0.1 or 0 for stronger voice capture
• Use Side-Step's Fisher Information analysis for adaptive per-layer rank allocation

Source: Voice cloning discussion (Issue #259), Side-Step toolkit

Test Coverage

25 tests cover ACE-Step-specific functionality:

• 7 detection tests: v1.0 (transformer_blocks, speaker_embedder, lyric_encoder) and v1.5 (PEFT, bare layers, cross_attn, MLP-only negative case)
• 12 normalization tests: v1.5 PEFT prefix stripping, bare layers prefix, Kohya underscore conversion, MLP keys, v1.0→v1.5 mapping (transformer_blocks→layers, attn→self_attn, to_q→q_proj, to_out.0→o_proj, cross_attn preserved), speaker_embedder, lyric_encoder (linear_q/v→q_proj/v_proj), no double-prefixing
• 6 preset tests: dedicated preset mapping, wider orthogonal band, higher TIES threshold, full magnitude preservation, auto-resolve, manual override

References

• ACE-Step 1.5 GitHub
• ACE-Step 1.5 Paper (arXiv:2602.00744)
• ACE-Step v1.0 Paper (arXiv:2506.00045)
• ACE-Step 1.5 config.json
• ACE-Step 1.5 Model Code
• ACE-Step 1.5 Explained (Substack)
• ACE-Step Architecture Analysis (zenn.dev)
• Voice Cloning Issue #259
• Side-Step Training Toolkit
• Community Merged LoRAs (DisturbingTheField)
• Community Merged LoRAs (acoustic guitar)
• LoRA Training Tutorial
• Musician's Guide Discussion
• LoRA Discussion/Library