Strong OOD generation but overfitting on the in-domain prompt pattern

When KL divergence is minimized, the student model is expected to generate images that are nearly identical to those of the teacher model. For example, with a GenEval-style prompt where both "a photo of" and "on the left of" appear together, the student model replicates the exact style of the GenEval-Teacher:

a photo of a cat on the left of a sign saying "hello world"

However, modifying either "a photo of" or "on the left of" pushes the prompt out of this overfitting region:

a cat on on the left of a sign saying "hello world"

a photo of a cat on holding a sign saying "hello world"

I assume you have observed this behavior, as the images presented in your paper are highly impressive and exhibit strong OOD generalization (like the two good images above), distinct from what is in train.txt and test.txt. This OOD improvement indeed demonstrates the effectiveness of the proposed OPD algorithm. However, the in-domain degradation toward each specific teacher model still persists.

Similarly, selecting any prompt from ocr/test.txt also results in this mode collapse/degradation toward the teacher's distribution. For example, for the first prompt in ocr/text.txt:

A high-fashion runway with a sleek, modern backdrop displaying "Spring Collection 2024". Models walk confidently on the catwalk, showcasing vibrant, floral prints and pastel tones, under soft, ambient lighting that enhances the fresh, spring vibe.

This behavior seems reasonable to me, given that the student model learns via step-wise KL divergence. This naturally forces its outputs to closely match the teacher's distribution within certain domains (e.g., train.txt and test.txt), inevitably inheriting the same limitations or degradation as the domain-specific teacher.

I am curious if you have considered methods to mitigate this overfitting/memorization on the training prompt set? I would love to hear your insights or thoughts on this.

Thank you again for this excellent and inspiring work! It has given me a lot to think about.

This is my inference code:

import torch
from diffusers import StableDiffusion3Pipeline
from peft import PeftModel

# 1. Load pipeline
pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium",
    torch_dtype=torch.bfloat16
)

# 2. Load LoRA weights
checkpoint = 'quanhaol/DiffusionOPD'
subfolder = 'Student/lora'
pipe.transformer = PeftModel.from_pretrained(
    pipe.transformer,
    checkpoint,
    subfolder=subfolder,
    torch_dtype=torch.bfloat16
)

pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power

# 3. Generate image
# prompt = "a photo of a cat holding a sign that says \"hello world\""
# prompt = "a photo of a cat on the left of a sign saying \"hello world\"" # GenEval-style, with `a photo of` and `on the left of`
# prompt = "a photo of a cat on holding a sign saying \"hello world\"" # with `a photo of` only
# prompt = "a cat on on the left of a sign saying \"hello world\"" # with `on the left of` only
prompt = "A high-fashion runway with a sleek, modern backdrop displaying \"Spring Collection 2024\". Models walk confidently on the catwalk, showcasing vibrant, floral prints and pastel tones, under soft, ambient lighting that enhances the fresh, spring vibe."

image = pipe(
    prompt,
    height=512,
    width=512,
    guidance_scale=4.5,
    num_inference_steps=40,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("output.png")

Start by reproducing the reported behavior with the provided StableDiffusion3Pipeline inference code and fixed seed, comparing prompts from train.txt, test.txt, and ocr/test.txt. Read the OPD training or distillation entry points referenced by the repository and evaluate whether a mitigation reduces in-domain overfitting without losing the reported OOD behavior; done requires an agreed method and supporting comparisons.