LaionBox v0.1 — Differentiable-Reward LoRA for DramaBox TTS
A LoRA adapter for the DramaBox LTX-2.3 text-to-speech model, trained with differentiable auxiliary rewards (CLAP naturalness, centroid real/fake scoring, WavLM speaker similarity) to improve voice naturalness while preserving speaker identity.
What This Is
This is a rank-128 LoRA (226M parameters, 906MB) that modifies the audio self-attention and feed-forward layers of the LTX-2.3 22B transformer. It was trained on 3,845 samples (3,247 DramaBox synthetic + 598 Emilia real speech) for 19 epochs total:
Epochs 1–12: Base IC-LoRA training (voice cloning with flow matching loss only)
Epochs 13–14: + CLAP auxiliary loss (positive/negative text similarity)
Epochs 15–19: + Differentiable reward training with 3 auxiliary losses through frozen reward models
The LoRA targets 6 module types across 10 transformer blocks:
The model produces audio that scores significantly higher on CLAP-based naturalness metrics — sounding less robotic and more like genuine human speech — while maintaining speaker identity fidelity above 0.89 cosine similarity.
An A100 80GB or equivalent GPU (the base model is 22B parameters)
Inference
python
1# From the DramaBox repository root:2python src/inference.py \
3--voice-sample reference.wav \
4--prompt "(softly, with warmth) \"I've been thinking about what you said.\"" \
5--checkpoint models/ltx-2.3-22b-dev-audio-only-v13-merged.safetensors \
6--full-checkpoint models/ltx-2.3-22b-dev.safetensors \
7--lora path/to/lora_epoch5.safetensors \
8--lora-rank 128 \
9--output output.wav
The --lora flag loads the LoRA weights, applies them via PEFT, and merges into the base model before inference. The inference script auto-detects the PEFT weight format and handles key mapping.
1from peft import LoraConfig, get_peft_model
2from safetensors.torch import load_file
34# After building the LTX velocity model:5lora_config = LoraConfig(6 r=128,7 lora_alpha=128,8 lora_dropout=0.0,9 bias="none",10 target_modules=[11"audio_attn1.to_k","audio_attn1.to_q","audio_attn1.to_v",12"audio_attn1.to_out.0","audio_ff.net.0.proj","audio_ff.net.2",13],14)15velocity_model = get_peft_model(velocity_model, lora_config)1617lora_sd = load_file("lora_epoch5.safetensors")18# Map key format if needed (lora_A.weight -> lora_A.default.weight)19mapped ={}20for k, v in lora_sd.items():21 k = k.replace(".lora_A.weight",".lora_A.default.weight")22 k = k.replace(".lora_B.weight",".lora_B.default.weight")23 mapped[k]= v
24velocity_model.load_state_dict(mapped, strict=False)25velocity_model = velocity_model.merge_and_unload()
Prompt Format
DramaBox uses a specific prompt format with stage directions in parentheses and dialogue in double quotes:
A young woman with a warm, slightly breathy voice delivers this studio-quality recording.
(gently, almost whispering) "I never thought I'd see you again."
(pause, then with growing emotion) "After all these years..."
For multi-scene prompts with emotional transitions:
A confident male speaker with a deep baritone delivers this high-quality studio recording.
(calmly, measured) "The reports all check out. Everything's in order."
CUT TO:
(urgently, voice cracking) "We need to evacuate. Now. Right now."
Training Data
The training dataset consists of 3,845 samples:
3,247 DramaBox samples: Synthetic TTS audio generated by the base DramaBox model, with MOSS-Audio-refined prompts that match the actual performance (not the original generation prompt)
598 Emilia samples: Real human speech from the Emilia dataset, filtered to the top 10% by DNS MOS score (≥ 3.478)
Each sample includes:
Pre-encoded audio latents (DramaBox VAE, 8×T×16 at ~25fps)
Pre-encoded text conditions (Gemma-3-12B-IT 4-bit hidden states)
Training mode labels: voice_clone_fwd, voice_clone_rev, unconditional
This was computed inside torch.no_grad() — the reward didn't backpropagate through the decoder or CLAP model. It provided a weak signal (rewards stayed flat). Best flow loss: 0.4474.
Differentiable Reward Training (Epochs 15–19) — The Key Innovation
The critical insight: reward-weighted reconstruction loss always produces gradients in the same direction as flow matching (toward x0_clean), regardless of the reward value. The reward only scales the magnitude, never the direction. To actually steer the model toward higher-reward outputs, you need gradients through the reward function itself.
We removed torch.no_grad() from the reward computation path, allowing gradients to flow:
VoiceCLAP's compute_log_mel has @torch.no_grad(): Despite torch.stft being fully differentiable, the decorator kills all gradients. We wrote encode_clap_waveform_differentiable() that replicates the mel computation without the decorator.
Wav2Vec2ForXVector loads wrong weights for WavLM: The checkpoint stores keys as wavlm.encoder.layers.* but Wav2Vec2ForXVector expects wav2vec2.encoder.layers.*. Fix: use WavLMForXVector.
With target_ratio=5.0 and coeff_cap=10.0, each aux loss targets 5× the flow matching magnitude.
Sigma Threshold
Auxiliary losses are only computed when the diffusion noise level sigma < 0.4. At high sigma, the x0 prediction is dominated by noise, so decoded audio is meaningless — reward signals would be random. About 23% of micro-batches pass this threshold.
Training Runs Summary
We ran 7 iterations to arrive at the final approach:
Audio latent: Encoded through DramaBox's VAE (8×T×16 at ~25fps)
Text condition: Gemma-3-12B-IT hidden states for the prompt
Mode: voice_clone_fwd, voice_clone_rev, or unconditional
Reference latent (optional): For voice cloning modes
See training_code/dramabox_finetune_train_multi_aux.py for the full data loading pipeline.
Key Hyperparameters
Parameter
Our value
Notes
lora_rank
128
Higher = more capacity, more VRAM
lora_alpha
128
Equal to rank (standard)
lr
4e-5
Peak learning rate
grad_accum
32
Global batch = GPUs × grad_accum
aux_target_ratio
5.0
Each aux loss targets 5× flow magnitude
coeff_cap
10.0
Max coefficient multiplier
aux_sigma_max
0.4
Only compute aux when sigma < this
differentiable_reward
true
Backprop through frozen reward models
diff_reward_checkpoint
true
Gradient checkpointing on aux models
VRAM Requirements
Component
VRAM
LTX-2.3 transformer (bf16, grad ckpt)
~44 GB
LoRA (rank 128) + AdamW optimizer
~3.7 GB
AudioDecoder + BigVGAN vocoder (frozen)
~0.5 GB
VoiceCLAP-small (frozen)
~0.3 GB
WavLM-SV (frozen, float32)
~0.2 GB
Quality MLP + centroids
~0.01 GB
Activations (checkpointed)
~18-25 GB
Total
~67-72 GB
Requires A100 80GB or equivalent. H100 80GB also works.
Insights: What We Learned About Auxiliary Losses for Diffusion Models
Reward-weighted reconstruction does not work
If you compute a scalar reward R and multiply it by the x0 reconstruction loss, the gradient is R * ∇(||x0_pred - x0_clean||²). This always points toward x0_clean — the reward R only changes the step size, never the direction. No matter how sophisticated your reward function, the model learns the same thing as pure flow matching, just faster or slower.
The entire path from trainable parameters to the loss function must allow gradient flow. A single @torch.no_grad() decorator anywhere in the chain (even on a seemingly internal helper function like mel spectrogram computation) will silently kill all gradient information. The training will appear to work (loss values exist, no errors) but rewards will be flat.
VoiceCLAP's mel computation kills gradients
VoiceCLAP-small (and likely other audio CLAP models) decorates compute_log_mel with @torch.no_grad(). The torch.stft inside is fully differentiable — the decorator is unnecessary and harmful for reward training. Our encode_clap_waveform_differentiable() function replicates the mel computation without the decorator, enabling proper gradient flow.
Sigma threshold is essential
At high noise levels (sigma > 0.4), the x0 prediction noisy - sigma * velocity_pred is dominated by noise. Decoding this through the VAE produces meaningless audio, and reward signals computed on it are random noise. Only compute auxiliary losses in the low-sigma regime where predictions are close to clean audio.
Centroid loss may conflict with naturalness
The centroid real/fake score always hit its coefficient cap (10.0) and never improved, while naturalness steadily increased. These two objectives may be in tension — the "real" direction in CLAP space (toward Emilia training data distribution) is not identical to the "natural-sounding" direction (toward positive text prompts).
Speaker similarity is already near ceiling
WavLM-SV speaker similarity started at 0.904 and stayed within 0.89-0.91 throughout training. The IC-LoRA architecture already preserves speaker identity well. Differentiable speaker loss provides minimal additional benefit but doesn't hurt.
Quality probability trends upward reliably
The quality MLP probability (clean vs distorted detector) showed the clearest monotonic improvement: 0.873 → 0.939 over 5 epochs. This suggests the differentiable CLAP path successfully steers audio away from distortion patterns.
Limitations
Centroid score degraded slightly (-0.012 over 5 epochs). The model moved toward CLAP naturalness at the expense of embedding proximity to the real speech centroid. These objectives partially conflict.
Flow matching loss increased slightly (+0.012). This is the expected tradeoff when auxiliary losses steer the model — it sacrifices some distributional matching for perceptual quality.
Validation inference failed during training due to a missing bitsandbytes dependency (needed for Gemma 4-bit inference in the validation subprocess). The LoRA checkpoints themselves are fine.
Training data is small (3,845 samples). Larger datasets may yield stronger improvements.
Single naturalness objective. The positive/negative text prompts are hardcoded. Different prompt choices would steer quality in different directions.
DramaBox-specific. This LoRA only works with the LTX-2.3 22B audio-only transformer architecture.
