license: apache-2.0
base_model: black-forest-labs/FLUX.2-klein-4B-base
library_name: peft
tags:
- lora
- on-policy-distillation
- classifier-free-guidance
Reference-conditioned distillation — PDM student (λ = 2)
The Positive-Direction Matching student from the reference-conditioned half of our study of classifier-free guidance in on-policy distillation. A text-only student on FLUX.2-klein-4B-base is distilled along its own denoising trajectories from a teacher that additionally sees a reference image the student never gets.
pdm_lambda2/
student/ 31.5 MB the deliverable — text-only, no reference image at inference
teacher/ 31.5 MB the EMA teacher it was distilled from (decay 0.9999)
Both are LoRA adapters over a frozen black-forest-labs/FLUX.2-klein-4B-base transformer,
r = 64, α = 128, 15.7 M trainable parameters each, attention q/k/v/out + the add_* projections +
both feed-forward branches.
What this setting is for
This is the privileged-conditioning half of the paper's "when does NBA emerge?" comparison, and it is where naive CFG-composed matching breaks.
The teacher runs the same backbone in its image-editing interface: it receives the student-visited
noisy latent plus encoded reference-image latents, and appends a style instruction to the prompt.
For its negative prediction the text is replaced by the null prompt and the reference latents
are swapped for a different image — the per-record neg_ref in the training manifest. The
student runs the plain text-to-image interface and sees no image at all:
c⁺_S = y c⁻_S = ∅
c⁺_T = (y, r) c⁻_T = (∅, r⁻) r⁻ ≠ r
So both of the teacher's branches carry image information the student has no way to represent, and they carry different image information. That is the privileged-conditioning asymmetry in its sharpest form.
Under that asymmetry, matching only the CFG-composed velocity stops reducing both branch errors:
‖e₊‖ falls while ‖e₋‖ climbs. We call this Negative Branch Asymmetry. It is invisible at
the training guidance scale and surfaces as guidance-scale sensitivity at inference — style drift
and visual artifacts, worst at γ = 1 where CFG collapses to the positive prediction alone. PDM
supervises before composition and suppresses the negative-error growth; the student then tracks the
teacher across the sweep.
The contrasting benign setting — shared negative conditioning, where naive matching is fine — is
text-rendering/.
Training
| objective | PDM, --loss-mode pdm --pdm-lambda 2.0 |
| guidance | γ = 2, student rollout and teacher |
| sampler | native FLUX.2-klein scheduler, 30-step denoising grid |
| supervision | first K = 16 transitions of the rollout |
| teacher | EMA of the student, decay 0.9999, initialized from the same weights |
| optimiser | AdamW, constant lr 1e-4, β = (0.9, 0.999), wd 0, grad clip 1.0 |
| batch | 8 GPUs × 1 × 1 ⇒ 8 trajectories per update |
| steps | 2000 updates, seed 30 |
| precision | bf16 mixed precision, DeepSpeed ZeRO-2, gradient checkpointing |
| resolution | 512 × 512 |
| data | the 32-record effect_all set shipped with the code (21 held out), including its neg_ref column |
The teacher's style instruction, verbatim:
Follow the overall visual style of the reference image, including its color palette, lighting, texture, rendering technique, and artistic treatment.
The negative reference comes from the neg_ref field of each training record. Drop that column
and the teacher's negative branch falls back to reusing the positive reference — a milder
asymmetry, and not what this checkpoint was trained under.
The same launcher reproduces the other arms — only --loss-mode and --pdm-lambda differ:
LOSS_MODE=naive bash scripts/train_pdm.sh # ‖γe₊ + (1−γ)e₋‖²
LOSS_MODE=ibm bash scripts/train_pdm.sh # γ²‖e₊‖² + (γ−1)²‖e₋‖²
Use
The code lives in DOPSD-flux/ of the
paper's release. Point inference at the student directory:
huggingface-cli download Cuttle-fish-my/Rethinking-CFG-OPD-ckpts \
--include 'reference-conditioned/*' --local-dir ckpts
python scripts/infer.py \
--model black-forest-labs/FLUX.2-klein-4B-base \
--lora ckpts/reference-conditioned/pdm_lambda2/student \
--prompt "A young woman rests her cheek on one hand beside a table of pastries in a warm orange kitchen, rendered in [A] style." \
--steps 40 --height 512 --width 512 --guidance-scale 2.0 --output sample.png
Generate at 512 × 512, the training resolution — infer.py defaults to 1024 and the adapter
was never trained there.
Use the same prompt convention as training: the style is carried by a concept token, [A] through
[E], and a prompt without one falls back to base-model behaviour. Sweeping --guidance-scale
across 1.0 / 1.5 / 2.0 / 2.5 is what exposes the NBA effect when compared against a naive student.
teacher/ is included because it is what the student was chasing — reproducing the branch-error
diagnostic needs it. It is not usable through infer.py, which drives the text-to-image interface;
the teacher expects the editing interface with reference-image latents concatenated.
Licence
Apache-2.0. FLUX.2-klein-4B-base carries its own terms.