MUSE-AI-Lab/Concept-Brittleness

This is the research code and model for a paper on object-dependent concept brittleness: a text-to-image model can apply a style such as "ink sketch" to most objects but repeatedly miss a few, regardless of seed. The authors argue these failures are systematic blind spots rather than sampling noise, and they provide a pipeline that diagnoses and repairs them without retraining the backbone.

How it works

The mechanism is a sparse autoencoder trained separately at each denoising step. Classifier-free-guided denoising updates are collected on training prompts and one Top-K SAE is trained per step, making abstract style and attribute concepts separable in sparse space Research README. The authors filter unreliable reference samples with a threshold and average the rest into a class-level concept prototype per step.

At inference, an Interpolation SAE Module encodes the current update and the prototype, interpolates them with coefficient alpha, and decodes back. The correction is applied only during an early window of the trajectory; the paper's defaults are a 64-dimensional model, 256 latent width, k=10 and alpha=0.8 over the earliest 20% of steps Research README. Because intervention happens during sampling, the README reports almost no added inference time.

Use path and concrete workflow

Running the full pipeline means extracting step-wise updates, training the per-step SAE, selecting reference samples, building prototypes, and then sampling with the guided mode. The repository's scripts cover each stage, and generate.py also provides plain sampling and mean-ablation modes; alpha and prompts are set through its flags Research README.

A plausible fit is a researcher or practitioner who wants to check whether a backbone has consistent concept blind spots and to see whether prototype nudging closes some of them without fine-tuning the model.

Evaluation scope

The reported results are on held-out prompts at 256×256 with each backbone's official sampler settings and seed 2026 Research README. Five backbones are covered: SD 1.5, SD 3.5 Large, SDXL, PixArt-alpha and FLUX.1-dev. Gains are reported for style metrics (CLIP-I, CLIP-T, Style Alignment) and attribute metrics (BLIP-VQA Color/Texture/Shape, CLIPSim); this is a benchmark on those metrics and settings, not a general quality claim.

Requirements and limits

• Python 3.10 with PyTorch; experiments used one A800 80GB GPU, so backbone scale matters for reproducing them Research README.
• Model weights are not included. Trusted local copies go under models/ with the documented layout, or you pass explicit paths; each backbone has its own Hugging Face source weights layout.
• The correction trades off repair against layout preservation: larger alpha strengthens the requested style but can change the layout.
• The README does not declare a license for this repository, while the adapted backbone code in src/ remains under the original backbones' licenses Research README. Check those terms before redistributing or building on the code.
• Defaults are reported as paper settings; the README does not claim they transfer unchanged to every concept, resolution or backbone.