reference-conditioned/README.md

---
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/`](../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:

```bash
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/`](https://github.com/Cuttle-fish-my/Rethinking-CFG-OPD) of the
paper's release. Point inference at the **student** directory:

```bash
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.