| --- |
| license: openrail++ |
| base_model: stabilityai/stable-diffusion-xl-base-1.0 |
| tags: |
| - lora |
| - sdxl |
| - stable-diffusion-xl |
| - text-to-image |
| - cityscapes |
| - concept-aware-lora |
| - research |
| library_name: calora |
| pipeline_tag: text-to-image |
| --- |
| |
| # Improving CA-LoRA — head-sliced LoRA checkpoints |
|
|
| LoRA adapters from the reproduction-and-extension study of **CA-LoRA (Concept-Aware LoRA)** on |
| SDXL / Cityscapes. Every adapter here is a **head-sliced** LoRA: instead of adapting whole |
| attention projection matrices, it adapts only the ~2% of *individual attention heads* that a |
| concept-attribution (CA) measurement picked out at a chosen diffusion timestep. |
|
|
| These are the exact adapters that produced the numbers in Tables 3 and 5 of the report, and the |
| images in the companion dataset repo |
| [`chs35/improving-ca-lora-artifacts`](https://huggingface.co/datasets/chs35/improving-ca-lora-artifacts). |
|
|
| ## ⚠️ Custom format — not loadable with `peft` or `diffusers` |
|
|
| Off-the-shelf adapter libraries cannot express attachment to a *partial slice* of a weight matrix, |
| so these checkpoints use a custom layout. `PeftModel.from_pretrained(...)` and |
| `pipe.load_lora_weights(...)` will **not** work. |
|
|
| Each `.safetensors` file contains, per adapted projection: |
|
|
| | tensor | dtype | meaning | |
| | --- | --- | --- | |
| | `<module>.lora_A` | f32 | LoRA down-projection, shaped to the selected slice | |
| | `<module>.lora_B` | f32 | LoRA up-projection, shaped to the selected slice | |
| | `<module>.lora_head_indices` | **int64** | which attention heads of that projection this adapter covers — the source of truth at load time | |
|
|
| plus safetensors metadata: `granularity` (`head`), `axis` (`style` / `viewpoint`), `strategy`, |
| `lora_rank` (64), `lora_alpha` (64), `steps` (10000). |
|
|
| Slice direction follows the single definition in the code (`HEAD_SLICE_DIMS`): **type A** = |
| `attn.to_out.0`, sliced along its *input columns*; **type B** = `to_q` / `to_k` / `to_v`, sliced |
| along their *output rows*. Partial selection matrices are never silently widened to the full |
| matrix, and a load under a mismatched `granularity` or `axis` is refused rather than coerced. |
|
|
| ### Loading |
|
|
| Use the `calora` library from the **Improving-CA-LoRA** code repository: |
|
|
| ```bash |
| python scripts/generate_images.py \ |
| --lora <RUN>/lora-final.safetensors \ |
| --run-name <RUN> \ |
| --mode both |
| ``` |
|
|
| (`calora.train.lora.load_lora` is the underlying entry point.) |
|
|
| ## Runs, and the report rows they back |
|
|
| Each `<RUN>/` directory is one training arm. `_style` / `_viewpoint` is the **concept axis** the CA |
| measurement was taken on; the letter/timestep prefix is the **timestep selection criterion**. |
|
|
| | directory | selection criterion | axis | report Table 3 row | |
| | --- | --- | --- | --- | |
| | `A_t81_style` | t = 81 (the original paper's timestep) | style | "t = 81 (paper) / style" | |
| | `A_t81_viewpoint` | t = 81 (the original paper's timestep) | viewpoint | "t = 81 (paper) / viewpoint" | |
| | `B_top3_style` | multi-t, aggregated over the top-3 timesteps [41, 1, 81] | style | "multi-t [41, 1, 81] / style" | |
| | `B_top3_viewpoint` | multi-t, aggregated over the top-3 timesteps [41, 1, 81] | viewpoint | "multi-t [41, 1, 81] / viewpoint" | |
| | `C_t41_style` | t = 41 (this study's top-ranked timestep) | style | "t = 41 / style" | |
| | `C_t41_viewpoint` | t = 41 (this study's top-ranked timestep) | viewpoint | "t = 41 / viewpoint" | |
| | `D_t1_style` | t = 1 | style | "t = 1 / style" | |
| | `D_t1_viewpoint` | t = 1 | viewpoint | "t = 1 / viewpoint" | |
| | `E_random` | **control**: a random 2% of heads, no CA measurement | — | "random 2% (no CA)" | |
| | `E_all` | **control**: all attention projections, full LoRA, no CA selection | — | "full attention (no CA)" | |
|
|
| The same ten directories back the drift rows of **Table 5** (per-snapshot Jaccard / Spearman of the |
| selected set against the trained set). The report's third control, `control_base` (0%, base SDXL |
| with no adapter), needs no checkpoint and so has none here; its generated images are in the dataset |
| repo. |
|
|
| ## Checkpoints per run |
|
|
| Five files per run — intermediate checkpoints every 2,000 steps plus the final one at step 10,000: |
|
|
| ``` |
| <RUN>/lora-step002000.safetensors |
| <RUN>/lora-step004000.safetensors |
| <RUN>/lora-step006000.safetensors |
| <RUN>/lora-step008000.safetensors |
| <RUN>/lora-final.safetensors # step 10000 — the adapter evaluated in Table 3 |
| ``` |
|
|
| `lora-final.safetensors` is the one used for every reported metric; the intermediates exist so that |
| the CA-drift-during-training analysis can be re-run at other points on the trajectory. |
|
|
| ## `pilot_module_granularity/A_t81/` — granularity ablation |
| |
| An **undocumented earlier pilot run at *module* granularity** (whole projection matrices selected, |
| not individual heads), kept here as a granularity-ablation data point. It is *not* the source of any |
| number in the report — the reported `A_t81_*` runs are the head-granularity ones above. |
| |
| It differs in format as well as in granularity: it took the `peft` path (`granularity=module` is the |
| only mode `peft` supports here), so its keys are peft-style `...lora_A.default.weight`, it has no |
| `lora_head_indices` tensors, and its metadata predates the `granularity`/`axis` fields. 20 tensors |
| vs. the 168–1680 of the head-sliced runs; 34 MB for all five checkpoints. |
|
|
| ## Training setup |
|
|
| SDXL base 1.0 in **fp32**; Cityscapes `train` split (2,975 images) at 1024²; **10,000 steps**, |
| batch size 1, constant lr **1e-4** (AdamW, β = 0.9/0.999, weight decay 0.01, grad clip 1.0), LoRA |
| **rank 64 / alpha 64**, diffusion loss only, random crop + random horizontal flip, seed 0. Selection |
| budget: 2% of attention units by unit count. Full config is released as `config.yaml` in the code |
| repository's `results/sdxl_cityscapes/`. |
|
|
| ## Intended use and limitations |
|
|
| Research artifacts for reproducing and auditing the report — not production models. They are |
| fine-tuned on **Cityscapes** (German urban street scenes, daytime, vehicle-mounted camera), so they |
| absorb that domain narrowly and measurably lose prompt controllability relative to base SDXL; that |
| loss is the subject of the study, not a defect to work around. Anyone using these adapters must |
| respect the [Cityscapes dataset terms](https://www.cityscapes-dataset.com/license/) as well as the |
| SDXL [OpenRAIL++](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENSE.md) |
| use restrictions inherited from the base model. |
|
|
| ## Verifying integrity |
|
|
| `results/sha256_checkpoints.txt` in the code repository lists the sha256 of all 55 files here, |
| keyed by the same repo-relative paths. |
|
|
