Release Anima Style Embedding v2

.gitattributes

@@ -37,3 +37,8 @@ assets/generalization_and_gallery_size.png filter=lfs diff=lfs merge=lfs -text
 assets/learned_vs_frozen_baseline.png filter=lfs diff=lfs merge=lfs -text
 assets/reference_count_curve.png filter=lfs diff=lfs merge=lfs -text
 assets/training_validation_curves.png filter=lfs diff=lfs merge=lfs -text
+assets/architecture_imagegen_v1.png filter=lfs diff=lfs merge=lfs -text
+assets/domain_reference_count_imagegen_v1.png filter=lfs diff=lfs merge=lfs -text
+assets/factor_control_imagegen_v1.png filter=lfs diff=lfs merge=lfs -text
+assets/hero_shared_axis_imagegen_v1.png filter=lfs diff=lfs merge=lfs -text
+assets/set_retrieval_matrix_imagegen_v1.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,88 +1,140 @@
 ---
-library_name: pytorch
-pipeline_tag: feature-extraction
 license: other
 license_name: circlestone-labs-non-commercial-license
-license_link: LICENSE.md
-base_model:
-  - circlestone-labs/Anima
-  - timm/vit_so400m_patch16_siglip_gap_512.v2_webli
+license_link: https://huggingface.co/circlestone-labs/Anima/blob/main/LICENSE.md
+library_name: pytorch
+pipeline_tag: feature-extraction
+base_model: circlestone-labs/Anima
 tags:
   - anime
+  - illustration
   - style-embedding
   - image-retrieval
   - reference-image
-  - siglip2
+  - multi-reference
+  - dinov3
   - anima
-  - pytorch
+  - controllable-generation
 ---
 
-# Anima Style Embedding v1.0
+# Anima Style Embedding v2
 
-**A multi-reference style encoder for person-centric anime and illustration images.**
+A multi-reference shared-axis style encoder for anime and illustration images.
 
-Anima Style Embedding converts one to eight reference images into a normalized 512-dimensional style vector. It combines multi-level SigLIP2 features from the full image and face with internal Anima features, then grounds the result across the complete reference set. The embedding is designed for open-set style retrieval and reference-driven style transfer systems rather than a fixed bank of known artists.
+Anima Style Embedding v2 converts one to eight reference images into a normalized 1,024-dimensional embedding and a controllable style memory for Anima. It preserves spatial VAE evidence, adds DINOv3 visual context and Anima-aware descriptors, and represents Line, Color, Texture, and Layout as overlapping memberships on one 256-axis style field.
 
-| Model size | References | Embedding | Modalities |
-|---:|---:|---:|---|
-| 17.49M parameters | 1–8 images | 512-D, L2-normalized | Full image + face + Anima internals |
+| Model size | References | Embedding | Shared axes | Reference signals |
+|---|---:|---:|---:|---|
+| 51,991,809 parameters | 1-8 images | 1,024-D, L2-normalized | 256 | VAE + DINOv3-S+ + Anima |
 
-![Learned encoder versus frozen backbone](assets/learned_vs_frozen_baseline.png)
+![Anima Style Embedding v2 overview: 91.2% Top-1 and 97.4% Top-5 for a disjoint four-image Gallery and two-image Probe evaluation.](assets/hero_shared_axis_imagegen_v1.png)
 
 ## Highlights
 
-- **Open-set style representation:** unseen artists can be embedded directly from reference images.
-- **Native multi-reference grounding:** Reference Consensus Re-reading lets each image be interpreted in the context of the complete 1–8 image set.
-- **Full-image and face awareness:** the encoder combines global composition and rendering cues with face-specific style information.
-- **Anima-aware features:** three internal Anima blocks complement SigLIP2 visual features with representations aligned to the target generator.
-- **Strong retrieval gains:** unseen synthetic-style MRR improves from 32.0% with the frozen SigLIP2 Full+Face baseline to 60.8% with the learned encoder on the 4,500-way validation gallery.
+- **Style-disjoint evaluation:** retrieval is measured on 1,000 unseen style identities with separate Gallery and Probe image sets.
+- **Native multi-reference reasoning:** the encoder compares one to eight references before producing the style representation.
+- **One shared style field:** Line, Color, Texture, and Layout are overlapping memberships, not independent encoders.
+- **Controllable routing:** Auto, Assisted, and Manual modes support per-image factor weights plus a separate Face weight.
+- **Anima-ready memory:** the same shared axes produce the retrieval embedding and the style memory used by the native transfer path.
+- **Position-free Layout evidence:** phase-free spatial-power statistics capture composition-like tendencies without copying absolute object coordinates.
 
 ## Performance
 
-Retrieval uses cosine similarity between normalized query embeddings and five-image style prototypes.
+### Disjoint set-to-set retrieval
+
+Gallery and Probe embeddings are built from different images of the same style. The primary protocol uses cosine similarity, three fixed seeds, and a 1,000-style validation gallery. The sealed test split was not used for checkpoint selection or for the reported v2 results.
+
+The headline condition uses a four-image Gallery and a two-image Probe:
+
+| Top-1 | Top-5 |
+|---:|---:|
+| **91.23%** | **97.40%** |
+
+![Top-1 and Top-5 retrieval matrices for Gallery sets of 2, 4, and 8 images and Probe sets of 1, 2, 4, and 8 images.](assets/set_retrieval_matrix_imagegen_v1.png)
+
+Top-1:
 
-### 4,500-way validation — unseen identities
+| Gallery \ Probe | 1 | 2 | 4 | 8 |
+|---|---:|---:|---:|---:|
+| 2 references | 67.07% | 81.93% | 89.90% | 93.97% |
+| 4 references | 77.17% | 91.23% | 96.60% | 98.47% |
+| 8 references | 83.27% | 95.07% | 98.83% | 99.83% |
 
-The unseen subset contains 500 identities retrieved against a gallery of 4,500 styles.
+Top-5:
 
-| Source | Top-1 | Top-5 | MRR |
-|---|---:|---:|---:|
-| Synthetic | 51.8% | 71.7% | 60.8% |
-| Human | 56.4% | 75.3% | 65.0% |
+| Gallery \ Probe | 1 | 2 | 4 | 8 |
+|---|---:|---:|---:|---:|
+| 2 references | 84.73% | 93.27% | 96.97% | 98.23% |
+| 4 references | 91.23% | 97.40% | 99.27% | 99.70% |
+| 8 references | 94.03% | 98.60% | 99.80% | 100.00% |
 
-### 500-way sealed test — unseen identities
+Machine-readable Overall, Synthetic, and Human results are included in [`metrics/evaluation-set-matrix.json`](metrics/evaluation-set-matrix.json).
 
-| Source | Top-1 | Top-5 | MRR |
-|---|---:|---:|---:|
-| Synthetic | 73.8% | 90.3% | 81.1% |
-| Human | 76.1% | 92.0% | 83.0% |
+### Reference count and domain behavior
 
-The sealed test uses a smaller gallery than validation, so its absolute scores should be read independently.
+Human references have greater within-style variation when only one Probe image is available. With a two-image Gallery, the Synthetic/Human Top-1 gap narrows from 6.2 points at Probe-1 to 0.8 point at Probe-8.
 
-![Generalization and gallery-size sensitivity](assets/generalization_and_gallery_size.png)
+![Synthetic and Human Top-1 retrieval improve as the disjoint Probe reference set grows from one to eight images.](assets/domain_reference_count_imagegen_v1.png)
 
-### Reference count
+| Probe references | Synthetic Top-1 | Human Top-1 |
+|---:|---:|---:|
+| 1 | 69.20% | 63.00% |
+| 2 | 83.27% | 79.73% |
+| 4 | 90.73% | 88.47% |
+| 8 | 93.80% | 93.00% |
 
-More references consistently improve human-style grounding. Synthetic K=5 uses the locked controlled/matched/cross-content protocol and is shown separately from the generic K diagnostics.
+## Factor control and routing
 
-![Reference count evaluation](assets/reference_count_curve.png)
+Factor behavior is evaluated on 4,096 interventions from held-out transform families.
+
+![Held-out intervention accuracy, selective factor response, shared-axis overlap, and multi-reference routing results.](assets/factor_control_imagegen_v1.png)
+
+| Factor | Classification accuracy | Target / non-target response |
+|---|---:|---:|
+| Line | 88.48% | 1.761x |
+| Color | 64.16% | 1.406x |
+| Texture | 94.04% | 2.393x |
+| Layout | 91.70% | 1.970x |
+
+| Control metric | Result |
+|---|---:|
+| Balanced factor accuracy | 84.60% |
+| Strong > weak ordering | 98.83% of 2,048 pairs |
+| Median factor-outlier suppression | 3.643x |
+| Auto routing error reduction vs. uniform | 19.80% |
+| Axes materially shared by all four factors | 82 |
+| Effective rank | 248.35 / 256 |
+
+Attribute reconstruction reduces error by 31.2% relative to a train-mean predictor. Removing a factor membership raises its corresponding attribute error, confirming that the memberships are used by the prediction path rather than serving only as labels.
 
 ## Architecture
 
-The encoder projects full-image, face, and Anima tokens into a shared 512-wide space and reads them with four eight-head mixer blocks. Functional Style Factorization introduces four auxiliary style queries during training. Reference Consensus Re-reading computes a consensus across the reference set, re-reads every image under that context, and forms an agreement-weighted prototype.
+![Architecture from one to eight full and face references through VAE, DINOv3, and Anima signals into a shared 256-axis field, routing, a 1,024-D embedding, and Anima style memory.](assets/architecture_imagegen_v1.png)
+
+The production encoder has one spatial field and one aligned dictionary:
 
-| Input | Shape per batch |
+1. A convolutional stem preserves full-image and face VAE spatial fields.
+2. Each spatial token is projected into the same 256-axis dictionary.
+3. Per-axis moments and translation-invariant power statistics retain persistent and composition-like style evidence.
+4. DINOv3 and Anima projections add visual and generator-aware context to the aligned axes.
+5. A relation Transformer compares references before estimating per-axis reliability.
+6. Overlapping memberships expose Line, Color, Texture, and Layout views without splitting the representation into branches.
+7. The mixed axes produce both a normalized retrieval embedding and style memory for Anima conditioning.
+
+### Routing modes
+
+| Mode | Behavior |
 |---|---|
-| Full-image SigLIP2 features | `[B, R, 30, 1152]` |
-| Face SigLIP2 features | `[B, R, 30, 1152]` |
-| Face validity mask | `[B, R]` |
-| Anima internal features | `[B, R, 3, 4096]` |
-| Reference mask | `[B, R]` |
+| Auto | Uses learned per-reference, per-axis reliability. |
+| Assisted | Multiplies learned reliability by user-provided factor and Face weights. |
+| Manual | Uses user weights directly to choose factor suppliers. |
 
-SigLIP2 features use [`timm/vit_so400m_patch16_siglip_gap_512.v2_webli`](https://huggingface.co/timm/vit_so400m_patch16_siglip_gap_512.v2_webli) at 512×512 and three backbone depths. Anima features use 768×768 inputs and blocks 8, 18, and 26 at sigma 0.1.
+Weights normalize across references, so adding references does not automatically increase style strength. A zero Manual factor weight produces zero output contribution and zero gradient on that reference-factor path. Face controls face-token routing and is not treated as a fifth factor branch.
 
 ## Quick start
 
+The v2 checkpoint expects precomputed features. Raw-RGB VAE, DINOv3-S+, and Anima descriptor extraction is not bundled in this release candidate.
+
 ```python
 import json
 
@@ -90,59 +142,99 @@ import torch
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
-from modeling_anima_style_embedding import UnifiedStyleEncoder
+from modeling_anima_style_embedding import SharedAxisStyleEncoder
 
 repo_id = "Baragi-AI/Anima-Style-Embedding"
-config = json.load(open(hf_hub_download(repo_id, "config.json"), encoding="utf-8"))
-cfg = config["architecture"]
-
-encoder = UnifiedStyleEncoder(
-    width=cfg["width"],
-    blocks=cfg["blocks"],
-    heads=cfg["heads"],
-    embedding_dim=cfg["embedding_dim"],
-    dropout=cfg["dropout"],
-    use_anima=True,
-    reference_grounding=cfg["reference_grounding"],
-    functional_factors=cfg["functional_factors"],
-)
-encoder.load_state_dict(
-    load_file(hf_hub_download(repo_id, "model.safetensors")),
-    strict=True,
-)
+config_path = hf_hub_download(repo_id, "config.json")
+weights_path = hf_hub_download(repo_id, "model.safetensors")
+
+with open(config_path, encoding="utf-8") as file:
+    config = json.load(file)
+
+encoder = SharedAxisStyleEncoder(**config["architecture"])
+encoder.load_state_dict(load_file(weights_path), strict=True)
 encoder.eval()
-```
 
-```python
 with torch.inference_mode():
-    style_embedding, per_image_embeddings = encoder(
-        full_features,
-        reference_mask,
-        face_features=face_features,
-        face_mask=face_mask,
-        anima_features=anima_features,
+    output = encoder(
+        full_latents=full_latents,
+        face_latents=face_latents,
+        dino=dino_features,
+        anima=anima_features,
+        reference_valid=reference_valid,
+        face_valid=face_valid,
+        user_weights=user_weights,
+        mode="assisted",
     )
 
-assert style_embedding.shape == (full_features.shape[0], 512)
+style_embedding = output["embedding"]
+style_memory = output["style_memory"]
+factor_coordinates = output["factor_coordinates"]
+
+assert style_embedding.shape == (full_latents.shape[0], 1024)
 ```
 
-`modeling_anima_style_embedding.py` contains the matching encoder definition. The checkpoint expects precomputed SigLIP2 and Anima features; RGB and Anima feature extraction are supplied by downstream integrations.
+### Input contract
+
+| Input | Shape | Notes |
+|---|---|---|
+| Full VAE latents | `[B, R, 16, H, W]` | Spatial full-image evidence |
+| Face VAE latents | `[B, R, 16, Hf, Wf]` | Optional, masked by `face_valid` |
+| DINOv3 full+face summary | `[B, R, 6912]` | Cached DINOv3-S+ context |
+| Anima descriptors | `[B, R, 3, 4096]` | Three internal Anima blocks |
+| Reference validity | `[B, R]` | At least one valid reference per sample |
+| User weights | `[B, R, 5]` | Line, Color, Texture, Layout, Face |
+
+## Training data and split policy
+
+The base corpus contains 500,000 images:
+
+- 250,000 Anima-generated images from 5,000 Synthetic style identities;
+- 250,000 human-created illustrations from 5,000 Human style identities.
+
+Synthetic and Human identities remain separate even when artist names overlap. The split is identity-disjoint:
 
-## Training
+| Split | Images | Style identities | Use |
+|---|---:|---:|---|
+| Train | 400,000 | 8,000 | Optimization and training-only statistics |
+| Validation | 50,000 | 1,000 | Checkpoint selection and reported v2 metrics |
+| Test | 50,000 | 1,000 | Sealed and unused |
 
-The model was trained on 320,000 person-centric style records: 160,000 Anima-generated images and 160,000 human-created illustrations across 8,000 training identities. Synthetic and human examples are mixed within each batch. Training uses four RTX 5090 GPUs, BF16 DistributedDataParallel, 16-way episodes, two queries per prototype, variable 1–8 image support, global negatives, and weak hard-negative sampling. The 4,500-step checkpoint produced the best full-validation score.
+Factor training additionally uses 108,096 controlled intervention images. Training and validation intervention families are disjoint.
 
-![Training dynamics](assets/training_validation_curves.png)
+## Development findings
 
-The training record is available in [Weights & Biases](https://wandb.ai/1wndrla17-kyung-hee-university/anima-style-embedding/runs/5c9ae843a14f46529fe9aaa604a78d3b). Machine-readable evaluation results are included under `metrics/`.
+The final structure follows several measured corrections:
+
+- Global pooling and repeated low-dimensional compression removed local line and texture evidence. The accepted path preserves the VAE spatial field.
+- Independent factor branches duplicated capacity and forced artificial separation. Overlapping memberships on one dictionary retained shared evidence.
+- An earlier dense-coordinate implementation broke coordinate identity between training and production gating. Aligned per-axis pooling made coordinate `k` refer to the same axis everywhere.
+- FiLM improved retrieval, but updating it during factor tuning reduced retrieval. Retrieval-trained FiLM is kept and frozen during later stages.
+- Layout needed phase-free radial and directional power statistics. These improved Layout prediction without encoding absolute object locations.
+- Fixed validation panels, atomic checkpoints, domain-balanced episodes, and median per-record routing metrics prevented evaluation artifacts from selecting a model.
+
+See [`RESEARCH_REPORT.md`](RESEARCH_REPORT.md) for the full architecture, training stages, failed approaches, and acceptance evidence.
 
 ## Applications
 
 - reference-driven style transfer for Anima;
-- artist-style retrieval and clustering;
-- aggregation of multiple references into a single style prototype;
-- open-set style recognition for previously unseen artists;
-- similarity scoring and reference-set curation.
+- open-set style retrieval and clustering;
+- multi-reference style prototype construction;
+- per-reference Line, Color, Texture, Layout, and Face control;
+- reference-set curation and similarity scoring;
+- research on shared and controllable style representations.
+
+## Limitations
+
+- Retrieval and intervention metrics measure representation behavior, not generated-image fidelity.
+- Raw-RGB feature extraction is not bundled in this release candidate.
+- The sealed test split remains unused, so v2 does not claim sealed-test performance.
+- Color intervention selectivity is weaker than Line, Texture, and Layout.
+- Layout attribute causality is positive but weaker than the other attribute memberships.
+- Human single-reference retrieval remains harder than Synthetic retrieval.
+- Line, Color, Texture, and Layout are useful controls, not a complete or perfectly disentangled theory of style.
+- The model is optimized for anime and illustration images and has not been evaluated as a general photographic style encoder.
+- The embedding must not be used as evidence of authorship or identity.
 
 ## Files
 
@@ -151,29 +243,28 @@ The training record is available in [Weights & Biases](https://wandb.ai/1wndrla1
 | `model.safetensors` | Encoder weights |
 | `config.json` | Architecture and feature contract |
 | `modeling_anima_style_embedding.py` | PyTorch model definition |
-| `metrics/` | Training curves and retrieval results |
+| `metrics/` | Retrieval, factor, attribute, and routing evaluations |
+| `assets/` | Model-card figures |
+| `RESEARCH_REPORT.md` | Detailed research and implementation report |
 | `SHA256SUMS.json` | Release checksums |
-
-## Limitations
-
-- The model is optimized for person-centric anime and illustration images; other domains are not evaluated.
-- Retrieval metrics measure representation quality, not generated-image fidelity.
-- Style and artist-associated content preferences are not completely disentangled.
-- Anima features must follow the documented model, block, resolution, and sigma configuration.
+| `LICENSE.md` | CircleStone Labs Non-Commercial License v1.2 |
+| `CODE_LICENSE/` | License for original release code |
+| `NOTICE` | Required attribution notice |
+| `THIRD_PARTY_NOTICES.md` | Third-party licenses and attribution |
 
 ## License
 
-The model weights are distributed under the [CircleStone Labs Non-Commercial License v1.2](LICENSE.md). Commercial or production use requires appropriate authorization from CircleStone Labs.
+The model weights are a derivative of [circlestone-labs/Anima](https://huggingface.co/circlestone-labs/Anima) and are distributed under the [CircleStone Labs Non-Commercial License v1.2](https://huggingface.co/circlestone-labs/Anima/blob/main/LICENSE.md).
 
-The original source code supplied with this release is available under the [Apache License 2.0](CODE_LICENSE). Third-party components remain subject to their respective licenses; attribution details are listed in [NOTICE](NOTICE) and [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md).
+The license permits the model and its derivatives only for the non-commercial purposes defined by that license. Commercial or production use requires authorization from CircleStone Labs. The upstream license separately permits use of model outputs, subject to its terms and applicable law.
 
-Built on NVIDIA Cosmos.
+Original Baragi AI release code is provided under the Apache License 2.0. Third-party components remain under their respective licenses and are listed in `NOTICE` and `THIRD_PARTY_NOTICES.md`.
 
 ## Citation
 
 ```bibtex
-@misc{baragi2026animastyleembedding,
-  title        = {Anima Style Embedding v1.0},
+@misc{baragi2026animastyleembeddingv2,
+  title        = {Anima Style Embedding v2},
   author       = {{Baragi AI}},
   year         = {2026},
   howpublished = {Hugging Face model repository},

RESEARCH_REPORT.md

@@ -0,0 +1,275 @@
+# Anima Style Embedding v2: Research Report
+
+## Executive summary
+
+Anima Style Embedding v2 is a 52.0M-parameter, multi-reference style encoder for anime and illustration images. It maps one to eight reference images into a normalized 1,024-dimensional embedding and, at the same time, exposes a controllable shared style space for Line, Color, Texture, Layout, and Face-aware routing.
+
+The production checkpoint uses three complementary signals:
+
+- spatial Anima VAE latents, which retain local line, texture, and layout evidence;
+- DINOv3-S+ descriptors, which supply robust visual context;
+- three cached Anima internal descriptors, which align the representation with the target generator.
+
+All signals enter one aligned 256-axis style dictionary. Line, Color, Texture, and Layout are learned as overlapping memberships on that dictionary rather than independent branches. This preserves a common style representation while allowing factor-specific reference weighting and intervention training.
+
+The accepted Stage 4 checkpoint was selected on style-disjoint validation only. The sealed 1,000-style test split was not opened.
+
+## Released checkpoint
+
+| Item | Value |
+|---|---|
+| Architecture | `shared_axis_v3` |
+| Parameters | 51,991,809 |
+| Embedding | 1,024-D, L2-normalized |
+| Shared style axes | 256 |
+| Reference images | 1–8 |
+| Spatial resolution mode | `s12` |
+| Set reasoning | 2 layers, 8 heads |
+| Inputs | VAE full/face, DINOv3 full/face, Anima descriptors |
+| Factor controls | Line, Color, Texture, Layout |
+| Face control | Per-reference face weight |
+| Checkpoint SHA-256 | `79f1b467fa9fdb855c38ae91015371851ad574ca7a73c24b586b68701c37af14` |
+
+The release contains the encoder weights only. Optimizer state and private training metadata are not part of the public model.
+
+## Data and split policy
+
+The base corpus contains 500,000 images:
+
+- 250,000 Anima-generated images from 5,000 synthetic style identities;
+- 250,000 human-created illustrations from 5,000 human style identities.
+
+Synthetic and human identities remain distinct even when their artist names overlap. The split is identity-disjoint:
+
+| Split | Records | Style identities | Use |
+|---|---:|---:|---|
+| Train | 400,000 | 8,000 | optimization and normalization statistics |
+| Validation | 50,000 | 1,000 | checkpoint selection and reported v2 metrics |
+| Test | 50,000 | 1,000 | sealed; not used for v2 selection or reporting |
+
+Factor learning additionally uses 108,096 controlled intervention images. Training and validation use disjoint intervention families so that evaluation measures factor response rather than memorization of an image transform.
+
+## Architecture
+
+### Shared spatial field
+
+The encoder does not globally pool the VAE latent before style reasoning. A convolutional stem preserves a spatial field, and every spatial token is projected into the same 256-axis dictionary. Per-axis mean and log-standard-deviation statistics capture persistent appearance, while translation-invariant power statistics capture low- and mid-frequency energy, anisotropy, and diagonal structure without copying absolute object positions.
+
+### Context alignment
+
+DINOv3 and Anima descriptors condition the same aligned axes. The accepted model keeps the retrieval-trained FiLM transformation fixed during later intervention and routing stages because updating FiLM caused a measurable retrieval conflict. DINOv3 and Anima relation projections still provide per-reference context for set reasoning.
+
+### Overlapping factor memberships
+
+Line, Color, Texture, and Layout are sigmoid memberships over the shared dictionary. An axis may contribute to several factors. This is deliberate: line rendering, texture, and color handling are correlated in real styles, and forcing four disjoint branches discarded useful shared evidence.
+
+The accepted checkpoint retains:
+
+- 82 axes with material membership in every factor;
+- minimum factor coverage of 0.11045;
+- effective rank 248.35 out of 256.
+
+The decomposition is therefore controllable but not assumed to be complete. Style evidence outside the four named factors remains in the shared representation.
+
+### Multi-reference routing
+
+A relation Transformer compares all references before estimating per-axis reliability. The routing contract supports:
+
+- **Auto:** the model selects reliable references per axis;
+- **Assisted:** automatic reliability is multiplied by user-provided per-image factor weights;
+- **Manual:** user weights directly choose factor suppliers.
+
+Reference weights are normalized, so adding references does not automatically increase style strength. A zero user weight produces exactly zero contribution and zero gradient for that reference-factor path. Face weights control face-token routing without being added as a fifth factor coordinate.
+
+## Evaluation protocol
+
+The primary retrieval evaluation uses disjoint gallery and probe image sets:
+
+- gallery references: 2, 4, or 8 images per style;
+- probe references: 1, 2, 4, or 8 images per style;
+- cosine similarity between L2-normalized set embeddings;
+- three fixed seeds: 20260724, 20260725, and 20260726;
+- Overall: 1,000-way retrieval;
+- Synthetic and Human: 500-way retrieval each.
+
+This protocol is stricter and more representative than comparing one image with another image from the same artist. It measures whether a small probe set can retrieve a style prototype built from a separate gallery set.
+
+### Overall set-to-set retrieval
+
+Top-1:
+
+| Gallery \ Probe | 1 | 2 | 4 | 8 |
+|---|---:|---:|---:|---:|
+| 2 references | 67.07% | 81.93% | 89.90% | 93.97% |
+| 4 references | 77.17% | 91.23% | 96.60% | 98.47% |
+| 8 references | 83.27% | 95.07% | 98.83% | 99.83% |
+
+Top-5:
+
+| Gallery \ Probe | 1 | 2 | 4 | 8 |
+|---|---:|---:|---:|---:|
+| 2 references | 84.73% | 93.27% | 96.97% | 98.23% |
+| 4 references | 91.23% | 97.40% | 99.27% | 99.70% |
+| 8 references | 94.03% | 98.60% | 99.80% | 100.00% |
+
+The complete Overall, Synthetic, and Human matrices are included with the release. Every primary Top-1 and Top-5 cell remains within two percentage points of the accepted translation-invariant Stage 2 teacher. The worst deltas were -1.00 percentage point Overall, -1.40 Synthetic, and -1.27 Human.
+
+### Single-query regression panel
+
+The secondary diagnostic compares one query image with style sets containing 1, 2, 4, or 8 references.
+
+| Domain | Metric | 1 | 2 | 4 | 8 |
+|---|---|---:|---:|---:|---:|
+| Overall | Top-1 | 50.47% | 67.20% | 77.63% | 82.20% |
+| Overall | Top-5 | 72.70% | 83.50% | 90.03% | 93.40% |
+| Synthetic | Top-1 | 56.20% | 69.33% | 80.60% | 85.93% |
+| Synthetic | Top-5 | 76.07% | 86.93% | 92.80% | 94.33% |
+| Human | Top-1 | 48.33% | 63.60% | 74.93% | 79.80% |
+| Human | Top-5 | 70.00% | 80.40% | 88.40% | 91.93% |
+
+Synthetic Top-1 at two references regressed by 2.20 percentage points from its Stage 2 teacher. It is retained as a disclosed secondary limitation; the primary disjoint set-to-set matrix passed its preservation gate.
+
+## Factor and routing evaluation
+
+Evaluation uses 4,096 interventions from families that were never used for training.
+
+| Metric | Result |
+|---|---:|
+| Balanced factor accuracy | 84.60% |
+| Median target/non-target response ratio | 1.862× |
+| Level-response correlation | 0.260 |
+| Strong > weak ordering | 98.83% of 2,048 pairs |
+
+Factor-specific results:
+
+| Factor | Classification accuracy | Target/non-target response |
+|---|---:|---:|
+| Line | 88.48% | 1.761× |
+| Color | 64.16% | 1.406× |
+| Texture | 94.04% | 2.393× |
+| Layout | 91.70% | 1.970× |
+
+Color is the weakest controlled factor and does not meet the earlier 1.5× response target. The release should not be interpreted as perfect factor disentanglement.
+
+Attribute reconstruction is evaluated against a train-mean predictor:
+
+| Metric | Result |
+|---|---:|
+| Model error | 0.23806 |
+| Mean-predictor error | 0.34610 |
+| Relative error | 0.68785 |
+| Error reduction | 31.2% |
+
+Removing each factor's membership raises its corresponding attribute error by 2.033× for Line, 2.431× for Color, 1.433× for Texture, and 1.118× for Layout. The factors are therefore causally used, although Layout remains the least selective attribute subspace.
+
+Routing evaluation shows:
+
+- 3.643× median per-record suppression of a factor-outlier reference;
+- 4.557× ratio of mean target and non-target suppression;
+- 19.80% lower coordinate error for Auto routing than uniform averaging.
+
+## Training stages
+
+### Stage 1 — global style retrieval
+
+The model first learned an open-set global style embedding with style-disjoint validation, balanced Synthetic/Human episodes, variable reference counts, supervised contrastive learning, set-query consistency, Sub-center ArcFace, and variance/covariance regularization. DINO-only and DINO+Anima candidates were matched at equal steps; the Anima-aware model gained 11.2 percentage points in Top-5@1 by step 1,000.
+
+### Stage 2 — aligned factor representation
+
+Attribute reconstruction and membership-causality objectives were added while preserving retrieval. A coordinate-identity defect in the first implementation was corrected by making production dictionary index `k` identical to factor coordinate `k`. The accepted aligned model then learned soft overlapping memberships and translation-invariant spatial-power statistics.
+
+### Stage 3 — intervention response
+
+Controlled transforms trained locality, strength ordering, non-target preservation, and global identity preservation. Membership logits use a separate high learning rate and zero weight decay; a prior optimizer wiring error had accidentally assigned them the 1e-4 context learning rate.
+
+### Stage 4 — routing and continual preservation
+
+Stage 4 alternated clean retrieval, attribute reconstruction, intervention replay, factor-outlier routing, and weak/strong replay. Similarity-distribution distillation from the frozen Stage 2 teacher limited forgetting. Memberships were frozen because joint membership updates reduced shared coverage without improving control. The selected step-2,000 checkpoint passed the corrected robust routing and retrieval gates.
+
+## What failed and what was learned
+
+### Excessive compression loses style evidence
+
+Early plans normalized multiple backbones into a small common vector and then tokenized that vector again. This discarded local line and texture evidence. The production path retains the VAE spatial field and uses a larger aligned axis space.
+
+### Separate factor branches were the wrong abstraction
+
+Independent Line, Color, Texture, and Layout branches encouraged artificial separation and duplicated representation capacity. Overlapping memberships on one dictionary preserve shared evidence while still allowing factor-specific control.
+
+### Reconstruction and intervention are different jobs
+
+The first attribute probe accepted no reconstruction coordinates because it required one descriptor to be both reconstructable and selectively changed by every intervention. Reconstruction now defines interpretable axes; intervention deltas guide direction, strength, and locality.
+
+### Some handcrafted descriptors were misleading
+
+Contour turning-angle features and layout contour-shape statistics were removed after poor held-out behavior. Blur/sharpen was removed from the Line intervention family because it behaved primarily as Texture. Descriptor groups are retained only when they improve held-out prediction.
+
+### Coordinate identity must be literal
+
+The most consequential defect mixed dense projected coordinates while production gating used raw dictionary indices. A coordinate named `k` did not refer to the same style axis in training and inference. Aligned per-axis pooling fixed the contract and made factor controls load-bearing.
+
+### Factor separation was not a loss-coefficient problem
+
+Absolute locality hinges, relative locality, and diversity coefficient screens could not overcome the coordinate defect or nearly identical memberships. The useful changes were structural: aligned coordinates, a real membership learning rate, and translation-invariant spatial statistics.
+
+### FiLM is useful for retrieval but unstable during intervention tuning
+
+Removing FiLM significantly reduced retrieval. Updating it during intervention training also reduced retrieval. The accepted design keeps the retrieval-trained DINO/Anima FiLM path and freezes it during Stages 2–4.
+
+### Layout needed phase-free spatial evidence
+
+Mean and variance were insufficient for composition-like style. Low/mid radial power, anisotropy, and diagonal energy improved Layout error and causality without encoding absolute object positions. The implementation derives its frequency basis from the actual field resolution after an early 12×12 assumption failed on 8×8 face fields.
+
+### Evaluation infrastructure can invalidate a good run
+
+Several non-model failures were corrected:
+
+- validation panels now replay a fixed seed rather than advancing the sampler epoch;
+- checkpoints use atomic replacement to avoid incomplete ZIP files;
+- hard-negative episodes are domain-balanced;
+- W&B run IDs and resume policies are explicit;
+- routing logs distinguish ratio-of-means from the robust median per-record ratio;
+- Stage 4 uses source-relative membership coverage instead of an obsolete absolute threshold.
+
+### Artist-tag escaping was valid data
+
+Backslashes before parentheses in Anima artist tags were initially suspected to be corruption. They are required to represent literal parentheses in ComfyUI syntax. Cache repair was limited to genuinely invalid escaping.
+
+## Strengths
+
+- Strong retrieval from one reference and near-saturated retrieval from larger disjoint sets.
+- Native 1–8 reference reasoning rather than post-hoc mean pooling.
+- One representation shared by retrieval, factor control, routing, and future style transfer.
+- Explicit per-image Line, Color, Texture, Layout, and Face control.
+- Translation-invariant layout evidence that avoids copying absolute object locations.
+- Validation on unseen style identities and unseen intervention families.
+- Exact-zero Manual routing contract suitable for production controls.
+- Synthetic and Human domains are reported separately.
+
+## Limitations
+
+- Metrics measure representation and control behavior, not generated-image fidelity.
+- The v2 release expects precomputed VAE, DINOv3, and Anima features; raw-RGB extraction is not bundled yet.
+- The sealed test split remains unused, so v2 does not claim test-set performance.
+- Color intervention separation is weaker than Line, Texture, and Layout.
+- Layout membership causality is positive but comparatively weak.
+- Human single-reference retrieval remains harder than Synthetic retrieval.
+- The four named factors are useful controls, not a complete or perfectly disentangled theory of style.
+- The model is optimized for anime and illustration images and has not been evaluated as a general photographic style encoder.
+
+## Intended use
+
+The encoder is intended for:
+
+- reference-driven Anima style transfer;
+- open-set style retrieval and clustering;
+- multi-reference style prototype construction;
+- reference-set curation and similarity scoring;
+- research on controllable style representations.
+
+It is not a raw-image end-to-end product and should not be presented as evidence of authorship or identity.
+
+## Licensing
+
+The public weights are a derivative of Anima and are distributed under the CircleStone Labs Non-Commercial License v1.2. Commercial or production use of the model requires appropriate authorization from CircleStone Labs. Original release code is provided under Apache License 2.0; third-party components retain their own licenses. Generated-image licensing is governed separately by the upstream Anima terms and the user's generation workflow.
+

SHA256SUMS.json

@@ -1,9 +1,21 @@
 {
-  "model.safetensors": "6cdf126152927f6e31dc856a8d10b8decee3bb20ce4c0ca449206db65d4633ec",
-  "config.json": "7323cbf6061896453a88cd0f6302831b96278a52eb8b0b3c1500d84552a1e68f",
-  "README.md": "3f5dff8ec693b6ae1587640983237b78a480bda94ffafd7087e15696cf55bb05",
-  "LICENSE.md": "0457e82ca9056fb71de47af655407fbddf135db22025bfcc6c540c21e0071c43",
   "CODE_LICENSE": "c71d239df91726fc519c6eb72d318ec65820627232b2f796219e87dcf35d0ab4",
+  "LICENSE.md": "0457e82ca9056fb71de47af655407fbddf135db22025bfcc6c540c21e0071c43",
   "NOTICE": "cc26915cc81a0c84c8d18294cb2a949454be5e17c509468c660008313cd0d59c",
-  "THIRD_PARTY_NOTICES.md": "82e8d6cfd59dd247a64beca52e45eba909737416316578dcd22273f1662cbb78"
+  "README.md": "07f16322497e0db68a130b5a3bf34a8e4feeee048942d7965327d9cacd76e98e",
+  "RESEARCH_REPORT.md": "2f6227ea78ab2851100052c34fb01f6499df0ae58aa9e799f0fecf27fdb1437b",
+  "THIRD_PARTY_NOTICES.md": "82e8d6cfd59dd247a64beca52e45eba909737416316578dcd22273f1662cbb78",
+  "__pycache__/modeling_anima_style_embedding.cpython-312.pyc": "a9fe32ff6aedd2c1be7f270caac2d2e9a94c6aef010d3a50388658f2c33a348b",
+  "assets/architecture_imagegen_v1.png": "8a99fe93f3974982b1d027a0a539ebba254b86a106f2a5f12475f90348c2aa05",
+  "assets/domain_reference_count_imagegen_v1.png": "c8258d4cdfb3d84967618605d75eab88e4d30a25c616cc10f0a0d51901260215",
+  "assets/factor_control_imagegen_v1.png": "0749a35eace0425dca0f4df8ee0550b64d911822a64700a8292aa87767cf8b67",
+  "assets/hero_shared_axis_imagegen_v1.png": "71810febfa8ce6683c4baff372d5854683b0266bb6d038d283698c8505ee360b",
+  "assets/set_retrieval_matrix_imagegen_v1.png": "a20f34d8cc0a3e19b5f72383347f09a70602bb10498cdbe5913744fe0ac5caf1",
+  "config.json": "7c2d00fa1b659f96c12044a3096a1ecb033395412e18f092f008d4889ce5dd18",
+  "metrics/evaluation-set-matrix.json": "dfbca1e313632a33644115b110330959036a01ed8cf92c137f7c2c8b8ea84d04",
+  "metrics/evaluation-single-query-3panel.json": "d35f42b065880857f2aac0e638d25bb211005d97317e2b86c996e3de982f6443",
+  "metrics/evaluation.json": "caacdb2cd8b4bc303b8ba5cc54799521b479b5d27a64fe9ff791a720fd6e04de",
+  "model.safetensors": "c7f38c64babe6a1dd711f7819654640b40102d6e73ae45cd2675dc84cda03453",
+  "modeling_anima_style_embedding.py": "bd20df4437a051abec62e9ad7d59ae242965cf39671b8e3984972457f2e39811",
+  "requirements.txt": "ac92c134bacf839bdd98dc18526d4dba2c1ce0e07c0a1b3f129b5a96ae4bfc70"
 }

config.json

@@ -1,75 +1,44 @@
 {
-  "model_type": "anima_style_embedding",
-  "version": "1.0",
+  "model_type": "anima_style_embedding_v2",
+  "architectures": [
+    "SharedAxisStyleEncoder"
+  ],
   "architecture": {
-    "width": 512,
-    "blocks": 4,
-    "heads": 8,
-    "embedding_dim": 512,
-    "dropout": 0.1,
-    "reference_grounding": true,
-    "functional_factors": true,
-    "reference_count": [
-      1,
-      8
-    ],
-    "trainable_parameters": 17488353
+    "dino_dim": 6912,
+    "style_dim": 1024,
+    "axis_count": 256,
+    "embedding_dim": 1024,
+    "anima_dim": 12288,
+    "resolution": "s12",
+    "set_layers": 2,
+    "set_heads": 8,
+    "direct_context_axes": true,
+    "axis_aligned": true,
+    "disable_token_film": false,
+    "translation_invariant_stats": true
   },
-  "inputs": {
-    "siglip2_full": {
-      "shape": [
-        30,
-        1152
-      ],
-      "dtype": "bfloat16"
-    },
-    "siglip2_face": {
-      "shape": [
-        30,
-        1152
-      ],
-      "dtype": "bfloat16",
-      "optional": true
-    },
-    "anima": {
-      "shape": [
-        3,
-        4096
-      ],
-      "dtype": "bfloat16"
-    }
-  },
-  "vision_backbone": {
-    "repo_id": "timm/vit_so400m_patch16_siglip_gap_512.v2_webli",
-    "image_size": 512,
-    "layer_fractions": [
-      0.25,
-      0.55,
-      1.0
-    ],
-    "pool_sizes": [
-      2,
-      4,
-      2
+  "feature_contract": {
+    "full_latents": "[B,R,16,H,W]",
+    "face_latents": "[B,R,16,Hf,Wf]",
+    "dino": "[B,R,6912]",
+    "anima": "[B,R,3,4096]",
+    "reference_valid": "[B,R]",
+    "face_valid": "[B,R]",
+    "user_weights": "[B,R,5]",
+    "factor_order": [
+      "line",
+      "color",
+      "texture",
+      "layout",
+      "face"
     ]
   },
-  "anima_features": {
-    "image_size": 768,
-    "blocks": [
-      8,
-      18,
-      26
-    ],
-    "sigma": 0.1,
-    "noise_seed": 20260715
-  },
-  "training": {
-    "checkpoint_step": 4500,
-    "face_dropout": 0.3,
-    "query_count": 2,
-    "optimizer": "AdamW(fused)",
-    "scheduler": "cosine_with_linear_warmup",
-    "learning_rate": 0.0001,
-    "packed_manifest_sha256": "c1267b9fd45c6e0d222b91d6ae2291b36b763277cc39cbc32484874ffeae9b32"
-  }
+  "routing_modes": [
+    "auto",
+    "assisted",
+    "manual"
+  ],
+  "parameter_count": 51991809,
+  "source_checkpoint_sha256": "79f1b467fa9fdb855c38ae91015371851ad574ca7a73c24b586b68701c37af14",
+  "validation_split": "style-disjoint validation; sealed test unused"
 }

metrics/evaluation-set-matrix.json

@@ -0,0 +1,341 @@
+{
+  "checkpoint": "/workspace/anima-style/runs/shared-axis-stage4-routing-focused/candidate-step2000.pt",
+  "domains": {
+    "human": {
+      "mean": {
+        "top1_gallery_r2_probe_r1": 0.6299999952316284,
+        "top1_gallery_r2_probe_r2": 0.7973333398501078,
+        "top1_gallery_r2_probe_r4": 0.8846666614214579,
+        "top1_gallery_r2_probe_r8": 0.9299999872843424,
+        "top1_gallery_r4_probe_r1": 0.7333333293596903,
+        "top1_gallery_r4_probe_r2": 0.8933333357175192,
+        "top1_gallery_r4_probe_r4": 0.9646666844685873,
+        "top1_gallery_r4_probe_r8": 0.9853333234786987,
+        "top1_gallery_r8_probe_r1": 0.8106666604677836,
+        "top1_gallery_r8_probe_r2": 0.9386666417121887,
+        "top1_gallery_r8_probe_r4": 0.9913333257039388,
+        "top1_gallery_r8_probe_r8": 0.9980000058809916,
+        "top5_gallery_r2_probe_r1": 0.8146666685740153,
+        "top5_gallery_r2_probe_r2": 0.9213333129882812,
+        "top5_gallery_r2_probe_r4": 0.9646666844685873,
+        "top5_gallery_r2_probe_r8": 0.9799999992052714,
+        "top5_gallery_r4_probe_r1": 0.8926666577657064,
+        "top5_gallery_r4_probe_r2": 0.972000002861023,
+        "top5_gallery_r4_probe_r4": 0.990666667620341,
+        "top5_gallery_r4_probe_r8": 0.9973333477973938,
+        "top5_gallery_r8_probe_r1": 0.9273333350817362,
+        "top5_gallery_r8_probe_r2": 0.984666665395101,
+        "top5_gallery_r8_probe_r4": 0.9973333477973938,
+        "top5_gallery_r8_probe_r8": 1.0
+      },
+      "panels": [
+        {
+          "seed": 20260724,
+          "top1_gallery_r2_probe_r1": 0.6539999842643738,
+          "top1_gallery_r2_probe_r2": 0.8259999752044678,
+          "top1_gallery_r2_probe_r4": 0.8939999938011169,
+          "top1_gallery_r2_probe_r8": 0.9340000152587891,
+          "top1_gallery_r4_probe_r1": 0.7540000081062317,
+          "top1_gallery_r4_probe_r2": 0.8939999938011169,
+          "top1_gallery_r4_probe_r4": 0.9620000123977661,
+          "top1_gallery_r4_probe_r8": 0.9819999933242798,
+          "top1_gallery_r8_probe_r1": 0.8199999928474426,
+          "top1_gallery_r8_probe_r2": 0.9279999732971191,
+          "top1_gallery_r8_probe_r4": 0.9900000095367432,
+          "top1_gallery_r8_probe_r8": 0.9959999918937683,
+          "top5_gallery_r2_probe_r1": 0.8379999995231628,
+          "top5_gallery_r2_probe_r2": 0.9279999732971191,
+          "top5_gallery_r2_probe_r4": 0.9660000205039978,
+          "top5_gallery_r2_probe_r8": 0.9819999933242798,
+          "top5_gallery_r4_probe_r1": 0.9179999828338623,
+          "top5_gallery_r4_probe_r2": 0.9700000286102295,
+          "top5_gallery_r4_probe_r4": 0.9879999756813049,
+          "top5_gallery_r4_probe_r8": 0.9940000176429749,
+          "top5_gallery_r8_probe_r1": 0.9259999990463257,
+          "top5_gallery_r8_probe_r2": 0.9779999852180481,
+          "top5_gallery_r8_probe_r4": 0.9980000257492065,
+          "top5_gallery_r8_probe_r8": 1.0
+        },
+        {
+          "seed": 20260725,
+          "top1_gallery_r2_probe_r1": 0.6320000290870667,
+          "top1_gallery_r2_probe_r2": 0.8040000200271606,
+          "top1_gallery_r2_probe_r4": 0.8859999775886536,
+          "top1_gallery_r2_probe_r8": 0.9419999718666077,
+          "top1_gallery_r4_probe_r1": 0.7059999704360962,
+          "top1_gallery_r4_probe_r2": 0.9020000100135803,
+          "top1_gallery_r4_probe_r4": 0.9660000205039978,
+          "top1_gallery_r4_probe_r8": 0.9860000014305115,
+          "top1_gallery_r8_probe_r1": 0.8199999928474426,
+          "top1_gallery_r8_probe_r2": 0.9419999718666077,
+          "top1_gallery_r8_probe_r4": 0.9879999756813049,
+          "top1_gallery_r8_probe_r8": 1.0,
+          "top5_gallery_r2_probe_r1": 0.8119999766349792,
+          "top5_gallery_r2_probe_r2": 0.921999990940094,
+          "top5_gallery_r2_probe_r4": 0.9620000123977661,
+          "top5_gallery_r2_probe_r8": 0.9800000190734863,
+          "top5_gallery_r4_probe_r1": 0.8920000195503235,
+          "top5_gallery_r4_probe_r2": 0.972000002861023,
+          "top5_gallery_r4_probe_r4": 0.9900000095367432,
+          "top5_gallery_r4_probe_r8": 1.0,
+          "top5_gallery_r8_probe_r1": 0.9359999895095825,
+          "top5_gallery_r8_probe_r2": 0.9900000095367432,
+          "top5_gallery_r8_probe_r4": 0.9959999918937683,
+          "top5_gallery_r8_probe_r8": 1.0
+        },
+        {
+          "seed": 20260726,
+          "top1_gallery_r2_probe_r1": 0.6039999723434448,
+          "top1_gallery_r2_probe_r2": 0.7620000243186951,
+          "top1_gallery_r2_probe_r4": 0.8740000128746033,
+          "top1_gallery_r2_probe_r8": 0.9139999747276306,
+          "top1_gallery_r4_probe_r1": 0.7400000095367432,
+          "top1_gallery_r4_probe_r2": 0.8840000033378601,
+          "top1_gallery_r4_probe_r4": 0.9660000205039978,
+          "top1_gallery_r4_probe_r8": 0.9879999756813049,
+          "top1_gallery_r8_probe_r1": 0.7919999957084656,
+          "top1_gallery_r8_probe_r2": 0.9459999799728394,
+          "top1_gallery_r8_probe_r4": 0.9959999918937683,
+          "top1_gallery_r8_probe_r8": 0.9980000257492065,
+          "top5_gallery_r2_probe_r1": 0.7940000295639038,
+          "top5_gallery_r2_probe_r2": 0.9139999747276306,
+          "top5_gallery_r2_probe_r4": 0.9660000205039978,
+          "top5_gallery_r2_probe_r8": 0.9779999852180481,
+          "top5_gallery_r4_probe_r1": 0.8679999709129333,
+          "top5_gallery_r4_probe_r2": 0.9739999771118164,
+          "top5_gallery_r4_probe_r4": 0.9940000176429749,
+          "top5_gallery_r4_probe_r8": 0.9980000257492065,
+          "top5_gallery_r8_probe_r1": 0.9200000166893005,
+          "top5_gallery_r8_probe_r2": 0.9860000014305115,
+          "top5_gallery_r8_probe_r4": 0.9980000257492065,
+          "top5_gallery_r8_probe_r8": 1.0
+        }
+      ],
+      "styles": 500
+    },
+    "overall": {
+      "mean": {
+        "top1_gallery_r2_probe_r1": 0.6706666747728983,
+        "top1_gallery_r2_probe_r2": 0.8193333347638448,
+        "top1_gallery_r2_probe_r4": 0.8990000089009603,
+        "top1_gallery_r2_probe_r8": 0.9396666685740153,
+        "top1_gallery_r4_probe_r1": 0.7716666658719381,
+        "top1_gallery_r4_probe_r2": 0.9123333295186361,
+        "top1_gallery_r4_probe_r4": 0.9660000006357828,
+        "top1_gallery_r4_probe_r8": 0.984666665395101,
+        "top1_gallery_r8_probe_r1": 0.8326666553815206,
+        "top1_gallery_r8_probe_r2": 0.9506666660308838,
+        "top1_gallery_r8_probe_r4": 0.9883333245913187,
+        "top1_gallery_r8_probe_r8": 0.9983333349227905,
+        "top5_gallery_r2_probe_r1": 0.8473333517710367,
+        "top5_gallery_r2_probe_r2": 0.9326666593551636,
+        "top5_gallery_r2_probe_r4": 0.9696666598320007,
+        "top5_gallery_r2_probe_r8": 0.9823333422342936,
+        "top5_gallery_r4_probe_r1": 0.9123333295186361,
+        "top5_gallery_r4_probe_r2": 0.9740000168482462,
+        "top5_gallery_r4_probe_r4": 0.9926666617393494,
+        "top5_gallery_r4_probe_r8": 0.996999979019165,
+        "top5_gallery_r8_probe_r1": 0.940333346525828,
+        "top5_gallery_r8_probe_r2": 0.9860000014305115,
+        "top5_gallery_r8_probe_r4": 0.9980000058809916,
+        "top5_gallery_r8_probe_r8": 1.0
+      },
+      "panels": [
+        {
+          "seed": 20260724,
+          "top1_gallery_r2_probe_r1": 0.6859999895095825,
+          "top1_gallery_r2_probe_r2": 0.828000009059906,
+          "top1_gallery_r2_probe_r4": 0.8980000019073486,
+          "top1_gallery_r2_probe_r8": 0.9409999847412109,
+          "top1_gallery_r4_probe_r1": 0.777999997138977,
+          "top1_gallery_r4_probe_r2": 0.9169999957084656,
+          "top1_gallery_r4_probe_r4": 0.9620000123977661,
+          "top1_gallery_r4_probe_r8": 0.9829999804496765,
+          "top1_gallery_r8_probe_r1": 0.8389999866485596,
+          "top1_gallery_r8_probe_r2": 0.9440000057220459,
+          "top1_gallery_r8_probe_r4": 0.9869999885559082,
+          "top1_gallery_r8_probe_r8": 0.996999979019165,
+          "top5_gallery_r2_probe_r1": 0.8640000224113464,
+          "top5_gallery_r2_probe_r2": 0.9390000104904175,
+          "top5_gallery_r2_probe_r4": 0.9729999899864197,
+          "top5_gallery_r2_probe_r8": 0.984000027179718,
+          "top5_gallery_r4_probe_r1": 0.9259999990463257,
+          "top5_gallery_r4_probe_r2": 0.9760000109672546,
+          "top5_gallery_r4_probe_r4": 0.9929999709129333,
+          "top5_gallery_r4_probe_r8": 0.996999979019165,
+          "top5_gallery_r8_probe_r1": 0.9380000233650208,
+          "top5_gallery_r8_probe_r2": 0.9860000014305115,
+          "top5_gallery_r8_probe_r4": 0.9990000128746033,
+          "top5_gallery_r8_probe_r8": 1.0
+        },
+        {
+          "seed": 20260725,
+          "top1_gallery_r2_probe_r1": 0.675000011920929,
+          "top1_gallery_r2_probe_r2": 0.8270000219345093,
+          "top1_gallery_r2_probe_r4": 0.906000018119812,
+          "top1_gallery_r2_probe_r8": 0.9440000057220459,
+          "top1_gallery_r4_probe_r1": 0.7689999938011169,
+          "top1_gallery_r4_probe_r2": 0.9150000214576721,
+          "top1_gallery_r4_probe_r4": 0.9679999947547913,
+          "top1_gallery_r4_probe_r8": 0.9860000014305115,
+          "top1_gallery_r8_probe_r1": 0.8349999785423279,
+          "top1_gallery_r8_probe_r2": 0.9509999752044678,
+          "top1_gallery_r8_probe_r4": 0.9860000014305115,
+          "top1_gallery_r8_probe_r8": 1.0,
+          "top5_gallery_r2_probe_r1": 0.8510000109672546,
+          "top5_gallery_r2_probe_r2": 0.9359999895095825,
+          "top5_gallery_r2_probe_r4": 0.968999981880188,
+          "top5_gallery_r2_probe_r8": 0.984000027179718,
+          "top5_gallery_r4_probe_r1": 0.9139999747276306,
+          "top5_gallery_r4_probe_r2": 0.9700000286102295,
+          "top5_gallery_r4_probe_r4": 0.9909999966621399,
+          "top5_gallery_r4_probe_r8": 0.996999979019165,
+          "top5_gallery_r8_probe_r1": 0.9449999928474426,
+          "top5_gallery_r8_probe_r2": 0.984000027179718,
+          "top5_gallery_r8_probe_r4": 0.9959999918937683,
+          "top5_gallery_r8_probe_r8": 1.0
+        },
+        {
+          "seed": 20260726,
+          "top1_gallery_r2_probe_r1": 0.6510000228881836,
+          "top1_gallery_r2_probe_r2": 0.8029999732971191,
+          "top1_gallery_r2_probe_r4": 0.8930000066757202,
+          "top1_gallery_r2_probe_r8": 0.9340000152587891,
+          "top1_gallery_r4_probe_r1": 0.7680000066757202,
+          "top1_gallery_r4_probe_r2": 0.9049999713897705,
+          "top1_gallery_r4_probe_r4": 0.9679999947547913,
+          "top1_gallery_r4_probe_r8": 0.9850000143051147,
+          "top1_gallery_r8_probe_r1": 0.8240000009536743,
+          "top1_gallery_r8_probe_r2": 0.9570000171661377,
+          "top1_gallery_r8_probe_r4": 0.9919999837875366,
+          "top1_gallery_r8_probe_r8": 0.9980000257492065,
+          "top5_gallery_r2_probe_r1": 0.8270000219345093,
+          "top5_gallery_r2_probe_r2": 0.9229999780654907,
+          "top5_gallery_r2_probe_r4": 0.9670000076293945,
+          "top5_gallery_r2_probe_r8": 0.9789999723434448,
+          "top5_gallery_r4_probe_r1": 0.8970000147819519,
+          "top5_gallery_r4_probe_r2": 0.9760000109672546,
+          "top5_gallery_r4_probe_r4": 0.9940000176429749,
+          "top5_gallery_r4_probe_r8": 0.996999979019165,
+          "top5_gallery_r8_probe_r1": 0.9380000233650208,
+          "top5_gallery_r8_probe_r2": 0.9879999756813049,
+          "top5_gallery_r8_probe_r4": 0.9990000128746033,
+          "top5_gallery_r8_probe_r8": 1.0
+        }
+      ],
+      "styles": 1000
+    },
+    "synthetic": {
+      "mean": {
+        "top1_gallery_r2_probe_r1": 0.6919999917348226,
+        "top1_gallery_r2_probe_r2": 0.8326666752497355,
+        "top1_gallery_r2_probe_r4": 0.9073333342870077,
+        "top1_gallery_r2_probe_r8": 0.9380000034968058,
+        "top1_gallery_r4_probe_r1": 0.7846666574478149,
+        "top1_gallery_r4_probe_r2": 0.9153333306312561,
+        "top1_gallery_r4_probe_r4": 0.9699999888737997,
+        "top1_gallery_r4_probe_r8": 0.9859999815622965,
+        "top1_gallery_r8_probe_r1": 0.8306666612625122,
+        "top1_gallery_r8_probe_r2": 0.9586666623751322,
+        "top1_gallery_r8_probe_r4": 0.990666667620341,
+        "top1_gallery_r8_probe_r8": 0.9973333477973938,
+        "top5_gallery_r2_probe_r1": 0.8686666687329611,
+        "top5_gallery_r2_probe_r2": 0.9419999917348226,
+        "top5_gallery_r2_probe_r4": 0.9773333271344503,
+        "top5_gallery_r2_probe_r8": 0.9899999896685282,
+        "top5_gallery_r4_probe_r1": 0.9153333107630411,
+        "top5_gallery_r4_probe_r2": 0.9760000109672546,
+        "top5_gallery_r4_probe_r4": 0.9933333396911621,
+        "top5_gallery_r4_probe_r8": 0.9980000058809916,
+        "top5_gallery_r8_probe_r1": 0.9386666814486185,
+        "top5_gallery_r8_probe_r2": 0.9893333315849304,
+        "top5_gallery_r8_probe_r4": 0.9986666838328043,
+        "top5_gallery_r8_probe_r8": 1.0
+      },
+      "panels": [
+        {
+          "seed": 20260724,
+          "top1_gallery_r2_probe_r1": 0.6840000152587891,
+          "top1_gallery_r2_probe_r2": 0.8320000171661377,
+          "top1_gallery_r2_probe_r4": 0.9139999747276306,
+          "top1_gallery_r2_probe_r8": 0.9319999814033508,
+          "top1_gallery_r4_probe_r1": 0.7879999876022339,
+          "top1_gallery_r4_probe_r2": 0.921999990940094,
+          "top1_gallery_r4_probe_r4": 0.9779999852180481,
+          "top1_gallery_r4_probe_r8": 0.9879999756813049,
+          "top1_gallery_r8_probe_r1": 0.8299999833106995,
+          "top1_gallery_r8_probe_r2": 0.9639999866485596,
+          "top1_gallery_r8_probe_r4": 0.9940000176429749,
+          "top1_gallery_r8_probe_r8": 0.9959999918937683,
+          "top5_gallery_r2_probe_r1": 0.8820000290870667,
+          "top5_gallery_r2_probe_r2": 0.9559999704360962,
+          "top5_gallery_r2_probe_r4": 0.9760000109672546,
+          "top5_gallery_r2_probe_r8": 0.9879999756813049,
+          "top5_gallery_r4_probe_r1": 0.9319999814033508,
+          "top5_gallery_r4_probe_r2": 0.9819999933242798,
+          "top5_gallery_r4_probe_r4": 0.9959999918937683,
+          "top5_gallery_r4_probe_r8": 1.0,
+          "top5_gallery_r8_probe_r1": 0.9440000057220459,
+          "top5_gallery_r8_probe_r2": 0.9919999837875366,
+          "top5_gallery_r8_probe_r4": 0.9980000257492065,
+          "top5_gallery_r8_probe_r8": 1.0
+        },
+        {
+          "seed": 20260725,
+          "top1_gallery_r2_probe_r1": 0.7059999704360962,
+          "top1_gallery_r2_probe_r2": 0.8320000171661377,
+          "top1_gallery_r2_probe_r4": 0.9020000100135803,
+          "top1_gallery_r2_probe_r8": 0.9440000057220459,
+          "top1_gallery_r4_probe_r1": 0.7799999713897705,
+          "top1_gallery_r4_probe_r2": 0.9100000262260437,
+          "top1_gallery_r4_probe_r4": 0.9679999947547913,
+          "top1_gallery_r4_probe_r8": 0.9879999756813049,
+          "top1_gallery_r8_probe_r1": 0.8399999737739563,
+          "top1_gallery_r8_probe_r2": 0.9660000205039978,
+          "top1_gallery_r8_probe_r4": 0.9900000095367432,
+          "top1_gallery_r8_probe_r8": 0.9980000257492065,
+          "top5_gallery_r2_probe_r1": 0.8560000061988831,
+          "top5_gallery_r2_probe_r2": 0.9399999976158142,
+          "top5_gallery_r2_probe_r4": 0.9739999771118164,
+          "top5_gallery_r2_probe_r8": 0.9900000095367432,
+          "top5_gallery_r4_probe_r1": 0.9139999747276306,
+          "top5_gallery_r4_probe_r2": 0.9760000109672546,
+          "top5_gallery_r4_probe_r4": 0.9940000176429749,
+          "top5_gallery_r4_probe_r8": 1.0,
+          "top5_gallery_r8_probe_r1": 0.9380000233650208,
+          "top5_gallery_r8_probe_r2": 0.9919999837875366,
+          "top5_gallery_r8_probe_r4": 1.0,
+          "top5_gallery_r8_probe_r8": 1.0
+        },
+        {
+          "seed": 20260726,
+          "top1_gallery_r2_probe_r1": 0.6859999895095825,
+          "top1_gallery_r2_probe_r2": 0.8339999914169312,
+          "top1_gallery_r2_probe_r4": 0.906000018119812,
+          "top1_gallery_r2_probe_r8": 0.9380000233650208,
+          "top1_gallery_r4_probe_r1": 0.7860000133514404,
+          "top1_gallery_r4_probe_r2": 0.9139999747276306,
+          "top1_gallery_r4_probe_r4": 0.9639999866485596,
+          "top1_gallery_r4_probe_r8": 0.9819999933242798,
+          "top1_gallery_r8_probe_r1": 0.8220000267028809,
+          "top1_gallery_r8_probe_r2": 0.9459999799728394,
+          "top1_gallery_r8_probe_r4": 0.9879999756813049,
+          "top1_gallery_r8_probe_r8": 0.9980000257492065,
+          "top5_gallery_r2_probe_r1": 0.8679999709129333,
+          "top5_gallery_r2_probe_r2": 0.9300000071525574,
+          "top5_gallery_r2_probe_r4": 0.9819999933242798,
+          "top5_gallery_r2_probe_r8": 0.9919999837875366,
+          "top5_gallery_r4_probe_r1": 0.8999999761581421,
+          "top5_gallery_r4_probe_r2": 0.9700000286102295,
+          "top5_gallery_r4_probe_r4": 0.9900000095367432,
+          "top5_gallery_r4_probe_r8": 0.9940000176429749,
+          "top5_gallery_r8_probe_r1": 0.9340000152587891,
+          "top5_gallery_r8_probe_r2": 0.984000027179718,
+          "top5_gallery_r8_probe_r4": 0.9980000257492065,
+          "top5_gallery_r8_probe_r8": 1.0
+        }
+      ],
+      "styles": 500
+    }
+  }
+}

metrics/evaluation-single-query-3panel.json

@@ -0,0 +1,243 @@
+{
+  "candidates": {
+    "direct256": {
+      "anima_features": true,
+      "checkpoint": "/workspace/anima-style/runs/shared-axis-stage4-routing-focused/candidate-step2000.pt",
+      "domains": {
+        "human": {
+          "mean": {
+            "axis_effective_rank": 248.3503875732422,
+            "axis_max_offdiag_cosine": 0.3716677725315094,
+            "top1_r1": 0.4833333194255829,
+            "top1_r2": 0.6359999974568685,
+            "top1_r4": 0.7493333419164022,
+            "top1_r8": 0.7980000178019205,
+            "top5_r1": 0.699999988079071,
+            "top5_r2": 0.8040000001589457,
+            "top5_r4": 0.8840000033378601,
+            "top5_r8": 0.9193333387374878
+          },
+          "panels": [
+            {
+              "axis_effective_rank": 248.3503875732422,
+              "axis_max_offdiag_cosine": 0.3716677725315094,
+              "seed": 20260724,
+              "top1_r1": 0.46399998664855957,
+              "top1_r2": 0.6359999775886536,
+              "top1_r4": 0.7440000176429749,
+              "top1_r8": 0.7900000214576721,
+              "top5_r1": 0.7139999866485596,
+              "top5_r2": 0.8159999847412109,
+              "top5_r4": 0.8920000195503235,
+              "top5_r8": 0.9240000247955322
+            },
+            {
+              "axis_effective_rank": 248.3503875732422,
+              "axis_max_offdiag_cosine": 0.3716677725315094,
+              "seed": 20260725,
+              "top1_r1": 0.48399999737739563,
+              "top1_r2": 0.6399999856948853,
+              "top1_r4": 0.7459999918937683,
+              "top1_r8": 0.7960000038146973,
+              "top5_r1": 0.6959999799728394,
+              "top5_r2": 0.7900000214576721,
+              "top5_r4": 0.8840000033378601,
+              "top5_r8": 0.921999990940094
+            },
+            {
+              "axis_effective_rank": 248.3503875732422,
+              "axis_max_offdiag_cosine": 0.3716677725315094,
+              "seed": 20260726,
+              "top1_r1": 0.5019999742507935,
+              "top1_r2": 0.6320000290870667,
+              "top1_r4": 0.7580000162124634,
+              "top1_r8": 0.8080000281333923,
+              "top5_r1": 0.6899999976158142,
+              "top5_r2": 0.8059999942779541,
+              "top5_r4": 0.8759999871253967,
+              "top5_r8": 0.9120000004768372
+            }
+          ],
+          "records": 25000,
+          "std": {
+            "axis_effective_rank": 0.0,
+            "axis_max_offdiag_cosine": 0.0,
+            "top1_r1": 0.015520590819103444,
+            "top1_r2": 0.003265968608957246,
+            "top1_r4": 0.006182416289578324,
+            "top1_r8": 0.007483319318916783,
+            "top5_r1": 0.010198036315239446,
+            "top5_r2": 0.010708236862135389,
+            "top5_r4": 0.006531985884842726,
+            "top5_r8": 0.005249344174247969
+          },
+          "styles": 500
+        },
+        "overall": {
+          "mean": {
+            "axis_effective_rank": 248.3503875732422,
+            "axis_max_offdiag_cosine": 0.3716677725315094,
+            "top1_r1": 0.5046666661898295,
+            "top1_r2": 0.671999990940094,
+            "top1_r4": 0.7763333320617676,
+            "top1_r8": 0.8220000068346659,
+            "top5_r1": 0.7269999980926514,
+            "top5_r2": 0.8349999984105428,
+            "top5_r4": 0.9003333250681559,
+            "top5_r8": 0.9339999953905741
+          },
+          "panels": [
+            {
+              "axis_effective_rank": 248.3503875732422,
+              "axis_max_offdiag_cosine": 0.3716677725315094,
+              "seed": 20260724,
+              "top1_r1": 0.4860000014305115,
+              "top1_r2": 0.6639999747276306,
+              "top1_r4": 0.7739999890327454,
+              "top1_r8": 0.8109999895095825,
+              "top5_r1": 0.7409999966621399,
+              "top5_r2": 0.8389999866485596,
+              "top5_r4": 0.8989999890327454,
+              "top5_r8": 0.9309999942779541
+            },
+            {
+              "axis_effective_rank": 248.3503875732422,
+              "axis_max_offdiag_cosine": 0.3716677725315094,
+              "seed": 20260725,
+              "top1_r1": 0.5130000114440918,
+              "top1_r2": 0.6779999732971191,
+              "top1_r4": 0.7739999890327454,
+              "top1_r8": 0.8220000267028809,
+              "top5_r1": 0.7250000238418579,
+              "top5_r2": 0.828000009059906,
+              "top5_r4": 0.8989999890327454,
+              "top5_r8": 0.9369999766349792
+            },
+            {
+              "axis_effective_rank": 248.3503875732422,
+              "axis_max_offdiag_cosine": 0.3716677725315094,
+              "seed": 20260726,
+              "top1_r1": 0.5149999856948853,
+              "top1_r2": 0.6740000247955322,
+              "top1_r4": 0.781000018119812,
+              "top1_r8": 0.8330000042915344,
+              "top5_r1": 0.7149999737739563,
+              "top5_r2": 0.8379999995231628,
+              "top5_r4": 0.902999997138977,
+              "top5_r8": 0.9340000152587891
+            }
+          ],
+          "records": 50000,
+          "std": {
+            "axis_effective_rank": 0.0,
+            "axis_max_offdiag_cosine": 0.0,
+            "top1_r1": 0.013224554288104472,
+            "top1_r2": 0.005887845760751241,
+            "top1_r4": 0.0032998453573119375,
+            "top1_r8": 0.008981468424922569,
+            "top5_r1": 0.010708259127091117,
+            "top5_r2": 0.004966546871788705,
+            "top5_r4": 0.0018856219044783918,
+            "top5_r8": 0.0024494825401091443
+          },
+          "styles": 1000
+        },
+        "synthetic": {
+          "mean": {
+            "axis_effective_rank": 248.3503875732422,
+            "axis_max_offdiag_cosine": 0.3716677725315094,
+            "top1_r1": 0.5619999965031942,
+            "top1_r2": 0.6933333476384481,
+            "top1_r4": 0.8059999942779541,
+            "top1_r8": 0.859333316485087,
+            "top5_r1": 0.7606666684150696,
+            "top5_r2": 0.8693333268165588,
+            "top5_r4": 0.9279999931653341,
+            "top5_r8": 0.9433333476384481
+          },
+          "panels": [
+            {
+              "axis_effective_rank": 248.3503875732422,
+              "axis_max_offdiag_cosine": 0.3716677725315094,
+              "seed": 20260724,
+              "top1_r1": 0.5640000104904175,
+              "top1_r2": 0.6980000138282776,
+              "top1_r4": 0.7979999780654907,
+              "top1_r8": 0.8619999885559082,
+              "top5_r1": 0.7540000081062317,
+              "top5_r2": 0.871999979019165,
+              "top5_r4": 0.9279999732971191,
+              "top5_r8": 0.9380000233650208
+            },
+            {
+              "axis_effective_rank": 248.3503875732422,
+              "axis_max_offdiag_cosine": 0.3716677725315094,
+              "seed": 20260725,
+              "top1_r1": 0.5479999780654907,
+              "top1_r2": 0.6460000276565552,
+              "top1_r4": 0.8040000200271606,
+              "top1_r8": 0.843999981880188,
+              "top5_r1": 0.7599999904632568,
+              "top5_r2": 0.8619999885559082,
+              "top5_r4": 0.9300000071525574,
+              "top5_r8": 0.9440000057220459
+            },
+            {
+              "axis_effective_rank": 248.3503875732422,
+              "axis_max_offdiag_cosine": 0.3716677725315094,
+              "seed": 20260726,
+              "top1_r1": 0.5740000009536743,
+              "top1_r2": 0.7360000014305115,
+              "top1_r4": 0.8159999847412109,
+              "top1_r8": 0.871999979019165,
+              "top5_r1": 0.7680000066757202,
+              "top5_r2": 0.8740000128746033,
+              "top5_r4": 0.9259999990463257,
+              "top5_r8": 0.9480000138282776
+            }
+          ],
+          "records": 25000,
+          "std": {
+            "axis_effective_rank": 0.0,
+            "axis_max_offdiag_cosine": 0.0,
+            "top1_r1": 0.010708262837900284,
+            "top1_r2": 0.0368902168587073,
+            "top1_r4": 0.007483314008929943,
+            "top1_r8": 0.011585430934170549,
+            "top5_r1": 0.005734883585275591,
+            "top5_r2": 0.005249344174247969,
+            "top5_r4": 0.0016329964712711142,
+            "top5_r8": 0.0041096047714775565
+          },
+          "styles": 500
+        }
+      },
+      "human_minus_synthetic": {
+        "top1_r1": -0.07866667707761132,
+        "top1_r2": -0.05733335018157959,
+        "top1_r4": -0.056666652361551884,
+        "top1_r8": -0.061333298683166504,
+        "top5_r8": -0.024000008900960323
+      },
+      "resolution": "s12",
+      "step": 2000
+    }
+  },
+  "contract": {
+    "nested_reference_prefix": true,
+    "reference_counts": [
+      1,
+      2,
+      4,
+      8
+    ],
+    "same_target_across_reference_counts": true,
+    "seeds": [
+      20260724,
+      20260725,
+      20260726
+    ],
+    "split": "validation",
+    "test_used": false
+  }
+}

metrics/evaluation.json

@@ -0,0 +1,117 @@
+{
+  "attribute/attribute_error": 0.23806460597734524,
+  "attribute/attribute_error/human/color": 0.2041477712397719,
+  "attribute/attribute_error/human/layout": 0.3487045900660114,
+  "attribute/attribute_error/human/line": 0.2568663033834128,
+  "attribute/attribute_error/human/texture": 0.306499986136275,
+  "attribute/attribute_error/overall/color": 0.17828673733810343,
+  "attribute/attribute_error/overall/layout": 0.29342790455345175,
+  "attribute/attribute_error/overall/line": 0.21477947311496154,
+  "attribute/attribute_error/overall/texture": 0.27130544221554836,
+  "attribute/attribute_error/synthetic/color": 0.1525209903662535,
+  "attribute/attribute_error/synthetic/layout": 0.2381604293611011,
+  "attribute/attribute_error/synthetic/line": 0.17284771567126797,
+  "attribute/attribute_error/synthetic/texture": 0.23624006975823775,
+  "attribute/attribute_ratio": 0.6878512990755349,
+  "attribute/factor_std": 0.01133992150425911,
+  "attribute/mean_predictor_error": 0.3460989407118975,
+  "attribute/mean_predictor_error/human/color": 0.42768885221590647,
+  "attribute/mean_predictor_error/human/layout": 0.4193881548407108,
+  "attribute/mean_predictor_error/human/line": 0.43766723310491573,
+  "attribute/mean_predictor_error/human/texture": 0.42246386204507497,
+  "attribute/mean_predictor_error/overall/color": 0.3307754093894214,
+  "attribute/mean_predictor_error/overall/layout": 0.33926805037765684,
+  "attribute/mean_predictor_error/overall/line": 0.359268975772589,
+  "attribute/mean_predictor_error/overall/texture": 0.3572357347621366,
+  "attribute/mean_predictor_error/synthetic/color": 0.23421905634921739,
+  "attribute/mean_predictor_error/synthetic/layout": 0.25916127857816873,
+  "attribute/mean_predictor_error/synthetic/line": 0.28115958259758983,
+  "attribute/mean_predictor_error/synthetic/texture": 0.2922470102932322,
+  "attribute/membership_ablation_error/human/color": 0.4833354100328614,
+  "attribute/membership_ablation_error/human/layout": 0.3872401288902689,
+  "attribute/membership_ablation_error/human/line": 0.4958906247126295,
+  "attribute/membership_ablation_error/human/texture": 0.43578456714301117,
+  "attribute/membership_ablation_error/overall/color": 0.43335724416979005,
+  "attribute/membership_ablation_error/overall/layout": 0.3280958065628183,
+  "attribute/membership_ablation_error/overall/line": 0.4365888942086092,
+  "attribute/membership_ablation_error/overall/texture": 0.3886331248763336,
+  "attribute/membership_ablation_error/synthetic/color": 0.3835632295569392,
+  "attribute/membership_ablation_error/synthetic/layout": 0.26896133447411424,
+  "attribute/membership_ablation_error/synthetic/line": 0.37750567111737304,
+  "attribute/membership_ablation_error/synthetic/texture": 0.3416547381412804,
+  "attribute/membership_ablation_ratio/color": 2.430675722939336,
+  "attribute/membership_ablation_ratio/layout": 1.118147938459109,
+  "attribute/membership_ablation_ratio/line": 2.0327310048615463,
+  "attribute/membership_ablation_ratio/texture": 1.4324560602347596,
+  "attribute/membership_coverage_min": 0.11045104265213013,
+  "attribute/membership_mean": 0.2886149287223816,
+  "attribute/membership_overlap_axes": 82,
+  "attribute/membership_pair_difference": 0.38394951820373535,
+  "attribute/membership_within_factor_std": 0.34480148553848267,
+  "attribute/records": 4096,
+  "intervention/exact_strength_pairs": 2048,
+  "intervention/factor_accuracy/color": 0.6416015625,
+  "intervention/factor_accuracy/layout": 0.9169921875,
+  "intervention/factor_accuracy/line": 0.884765625,
+  "intervention/factor_accuracy/texture": 0.9404296875,
+  "intervention/factor_balanced_accuracy": 0.845947265625,
+  "intervention/factor_ratio/color": 1.405653715133667,
+  "intervention/factor_ratio/layout": 1.9703501462936401,
+  "intervention/factor_ratio/line": 1.7611885070800781,
+  "intervention/factor_ratio/texture": 2.39322829246521,
+  "intervention/family_accuracy/channel_mixer": 0.6416015625,
+  "intervention/family_accuracy/edge_overlay": 0.884765625,
+  "intervention/family_accuracy/median_speckle": 0.9404296875,
+  "intervention/family_accuracy/shear": 0.9169921875,
+  "intervention/level_response_correlation": 0.26010704040527344,
+  "intervention/loss/attribute_delta": 0.13530989782884717,
+  "intervention/loss/diversity": 0.3046875,
+  "intervention/loss/identity": 0.20753358397632837,
+  "intervention/loss/locality": 1.076119126752019,
+  "intervention/loss/preservation": 0.0003530628382577561,
+  "intervention/records": 4096,
+  "intervention/strong_gt_weak_exact": 0.98828125,
+  "intervention/target_non_target_ratio_median": 1.861896276473999,
+  "retrieval/axis_effective_rank": 248.3503875732422,
+  "retrieval/axis_max_offdiag_cosine": 0.3716677725315094,
+  "retrieval/top1_r1": 0.48500001430511475,
+  "retrieval/top1_r2": 0.6629999876022339,
+  "retrieval/top1_r4": 0.7739999890327454,
+  "retrieval/top1_r8": 0.8119999766349792,
+  "retrieval/top5_r1": 0.7409999966621399,
+  "retrieval/top5_r2": 0.8389999866485596,
+  "retrieval/top5_r4": 0.8960000276565552,
+  "retrieval/top5_r8": 0.9309999942779541,
+  "routing/auto_coordinate_error": 0.00034273043274879456,
+  "routing/auto_improvement": 0.19799659329327857,
+  "routing/non_target_attention_drop": 0.008984970860183239,
+  "routing/routing": 0.012377707287669182,
+  "routing/suppression_ratio_mean": 4.55674592245753,
+  "routing/suppression_ratio_median": 3.6429529190063477,
+  "routing/target_attention_drop": 0.0409422293305397,
+  "routing/uniform_coordinate_error": 0.00042734286398626864,
+  "set_retrieval/top1_gallery_r2_probe_r1": 0.6859999895095825,
+  "set_retrieval/top1_gallery_r2_probe_r2": 0.828000009059906,
+  "set_retrieval/top1_gallery_r2_probe_r4": 0.8980000019073486,
+  "set_retrieval/top1_gallery_r2_probe_r8": 0.9409999847412109,
+  "set_retrieval/top1_gallery_r4_probe_r1": 0.777999997138977,
+  "set_retrieval/top1_gallery_r4_probe_r2": 0.9169999957084656,
+  "set_retrieval/top1_gallery_r4_probe_r4": 0.9620000123977661,
+  "set_retrieval/top1_gallery_r4_probe_r8": 0.9829999804496765,
+  "set_retrieval/top1_gallery_r8_probe_r1": 0.8389999866485596,
+  "set_retrieval/top1_gallery_r8_probe_r2": 0.9440000057220459,
+  "set_retrieval/top1_gallery_r8_probe_r4": 0.9869999885559082,
+  "set_retrieval/top1_gallery_r8_probe_r8": 0.996999979019165,
+  "set_retrieval/top5_gallery_r2_probe_r1": 0.8640000224113464,
+  "set_retrieval/top5_gallery_r2_probe_r2": 0.9390000104904175,
+  "set_retrieval/top5_gallery_r2_probe_r4": 0.9729999899864197,
+  "set_retrieval/top5_gallery_r2_probe_r8": 0.984000027179718,
+  "set_retrieval/top5_gallery_r4_probe_r1": 0.9259999990463257,
+  "set_retrieval/top5_gallery_r4_probe_r2": 0.9760000109672546,
+  "set_retrieval/top5_gallery_r4_probe_r4": 0.9929999709129333,
+  "set_retrieval/top5_gallery_r4_probe_r8": 0.996999979019165,
+  "set_retrieval/top5_gallery_r8_probe_r1": 0.9380000233650208,
+  "set_retrieval/top5_gallery_r8_probe_r2": 0.9860000014305115,
+  "set_retrieval/top5_gallery_r8_probe_r4": 0.9990000128746033,
+  "set_retrieval/top5_gallery_r8_probe_r8": 1.0
+}

metrics/model_card_metrics.json

@@ -1,169 +0,0 @@
-{
-  "selected_checkpoint": {
-    "checkpoint": "/workspace/runs/style_embedding_combined_lr1e4_clip10_5k/best.pt",
-    "checkpoint_sha256": "16345685b72c9bb018833ea9d42d18ebdaafa585b2a2eb136f267006c1de5663",
-    "step": 4500,
-    "validation_checkpoint_score": 0.6301507651805878
-  },
-  "validation_k5": {
-    "synthetic": {
-      "styles": 4500.0,
-      "queries": 22500.0,
-      "top1": 0.5625777840614319,
-      "top5": 0.7585333585739136,
-      "mrr": 0.6518997550010681,
-      "margin": -0.006614952813833952,
-      "positive_cosine": 0.7365614771842957,
-      "hardest_negative_cosine": 0.7431764006614685,
-      "verification_auroc": 0.5534452199935913,
-      "embedding_std": 0.04349411278963089,
-      "embedding_effective_rank": 24.163583755493164,
-      "seen_queries": 20000.0,
-      "seen_top1": 0.5681999921798706,
-      "seen_top5": 0.7637499570846558,
-      "seen_mrr": 0.6573799848556519,
-      "seen_margin": -0.004584707319736481,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.5175999999046326,
-      "unseen_top5": 0.7167999744415283,
-      "unseen_mrr": 0.6080579161643982,
-      "unseen_margin": -0.02285691164433956,
-      "matched_queries": 4500.0,
-      "matched_top1": 0.48955556750297546,
-      "matched_top5": 0.6902222037315369,
-      "matched_mrr": 0.5840909481048584,
-      "matched_margin": -0.030851511284708977,
-      "cross_queries": 18000.0,
-      "cross_top1": 0.5808333158493042,
-      "cross_top5": 0.7756111025810242,
-      "cross_mrr": 0.6688519716262817,
-      "cross_margin": -0.0005558129050768912
-    },
-    "human": {
-      "styles": 4500.0,
-      "queries": 22500.0,
-      "top1": 0.6493777632713318,
-      "top5": 0.8219555616378784,
-      "mrr": 0.7280733585357666,
-      "margin": 0.022814957424998283,
-      "positive_cosine": 0.7811680436134338,
-      "hardest_negative_cosine": 0.7583531141281128,
-      "verification_auroc": 0.6369090676307678,
-      "embedding_std": 0.043804656714200974,
-      "embedding_effective_rank": 27.383806228637695,
-      "seen_queries": 20000.0,
-      "seen_top1": 0.660099983215332,
-      "seen_top5": 0.830549955368042,
-      "seen_mrr": 0.7378283143043518,
-      "seen_margin": 0.026178810745477676,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.5636000037193298,
-      "unseen_top5": 0.7531999945640564,
-      "unseen_mrr": 0.6500336527824402,
-      "unseen_margin": -0.004095893818885088
-    }
-  },
-  "sealed_test_k5": {
-    "synthetic": {
-      "styles": 500.0,
-      "queries": 2500.0,
-      "top1": 0.7375999689102173,
-      "top5": 0.9031999707221985,
-      "mrr": 0.8114251494407654,
-      "margin": 0.06940444558858871,
-      "positive_cosine": 0.7258960008621216,
-      "hardest_negative_cosine": 0.6564915776252747,
-      "verification_auroc": 0.7113653421401978,
-      "embedding_std": 0.04344247281551361,
-      "embedding_effective_rank": 23.896039962768555,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.7375999689102173,
-      "unseen_top5": 0.9031999707221985,
-      "unseen_mrr": 0.8114251494407654,
-      "unseen_margin": 0.06940444558858871,
-      "matched_queries": 500.0,
-      "matched_top1": 0.6600000262260437,
-      "matched_top5": 0.8520000576972961,
-      "matched_mrr": 0.748103678226471,
-      "matched_margin": 0.045967720448970795,
-      "cross_queries": 2000.0,
-      "cross_top1": 0.7570000290870667,
-      "cross_top5": 0.9160000681877136,
-      "cross_mrr": 0.8272556662559509,
-      "cross_margin": 0.0752636268734932
-    },
-    "human": {
-      "styles": 500.0,
-      "queries": 2500.0,
-      "top1": 0.761199951171875,
-      "top5": 0.9199999570846558,
-      "mrr": 0.830402135848999,
-      "margin": 0.08207859843969345,
-      "positive_cosine": 0.7568042874336243,
-      "hardest_negative_cosine": 0.6747256517410278,
-      "verification_auroc": 0.7453367114067078,
-      "embedding_std": 0.04376835376024246,
-      "embedding_effective_rank": 26.91031265258789,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.761199951171875,
-      "unseen_top5": 0.9199999570846558,
-      "unseen_mrr": 0.830402135848999,
-      "unseen_margin": 0.08207859843969345
-    }
-  },
-  "frozen_full_face_validation_k5": {
-    "synthetic": {
-      "styles": 4500.0,
-      "queries": 22500.0,
-      "top1": 0.2617333233356476,
-      "top5": 0.3952000141143799,
-      "mrr": 0.3287946879863739,
-      "margin": -0.058847684413194656,
-      "positive_cosine": 0.3081073760986328,
-      "hardest_negative_cosine": 0.36695507168769836,
-      "verification_auroc": 0.3488680124282837,
-      "seen_queries": 20000.0,
-      "seen_top1": 0.2628999948501587,
-      "seen_top5": 0.3959999978542328,
-      "seen_mrr": 0.3298960030078888,
-      "seen_margin": -0.0586077980697155,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.2523999810218811,
-      "unseen_top5": 0.3887999951839447,
-      "unseen_mrr": 0.31998398900032043,
-      "unseen_margin": -0.06076674535870552,
-      "matched_queries": 4500.0,
-      "matched_top1": 0.10533333569765091,
-      "matched_top5": 0.1851111203432083,
-      "matched_mrr": 0.14898858964443207,
-      "matched_margin": -0.11573364585638046,
-      "cross_queries": 18000.0,
-      "cross_top1": 0.3008333444595337,
-      "cross_top5": 0.4477222263813019,
-      "cross_mrr": 0.37374621629714966,
-      "cross_margin": -0.04462619498372078
-    },
-    "human": {
-      "styles": 4500.0,
-      "queries": 22500.0,
-      "top1": 0.43657779693603516,
-      "top5": 0.5786222219467163,
-      "mrr": 0.5056873559951782,
-      "margin": -0.029125619679689407,
-      "positive_cosine": 0.4710022509098053,
-      "hardest_negative_cosine": 0.5001279711723328,
-      "verification_auroc": 0.4406580626964569,
-      "seen_queries": 20000.0,
-      "seen_top1": 0.4320499897003174,
-      "seen_top5": 0.5762999653816223,
-      "seen_mrr": 0.5021449327468872,
-      "seen_margin": -0.030033092945814133,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.47279998660087585,
-      "unseen_top5": 0.5971999764442444,
-      "unseen_mrr": 0.5340264439582825,
-      "unseen_margin": -0.021865803748369217
-    }
-  },
-  "protocol_note": "Synthetic K=5 uses the locked five-controlled support, one matched and four cross-content query protocol. Other K values use deterministic generic record ordering."
-}

metrics/retrieval_metrics.csv

@@ -1,21 +0,0 @@
-split,source,k,styles,queries,top1,top5,mrr,margin,seen_top1,seen_top5,seen_mrr,unseen_top1,unseen_top5,unseen_mrr,matched_mrr,cross_mrr,verification_auroc,embedding_effective_rank
-validation,synthetic,1,4500.0,40500.0,0.25738272070884705,0.4447901248931885,0.3484785854816437,-0.13694550096988678,0.26072221994400024,0.448972225189209,0.35212650895118713,0.23066666722297668,0.4113333523273468,0.3192950487136841,,,0.2625024914741516,24.163583755493164
-validation,synthetic,2,4500.0,36000.0,0.38269445300102234,0.5837222337722778,0.479046106338501,-0.0735325887799263,0.3878125250339508,0.5885000228881836,0.48388049006462097,0.34175002574920654,0.5455000400543213,0.4403711259365082,,,0.3855718970298767,24.163583755493164
-validation,synthetic,4,4500.0,27000.0,0.4913703501224518,0.6868888735771179,0.5827372670173645,-0.032036565244197845,0.4964166581630707,0.6924583315849304,0.5878510475158691,0.45100000500679016,0.6423333287239075,0.5418276190757751,,,0.4890711009502411,24.163583755493164
-validation,synthetic,5,4500.0,22500.0,0.5625777840614319,0.7585333585739136,0.6518997550010681,-0.006614952813833952,0.5681999921798706,0.7637499570846558,0.6573799848556519,0.5175999999046326,0.7167999744415283,0.6080579161643982,0.5840909481048584,0.6688519716262817,0.5534452199935913,24.163583755493164
-validation,synthetic,8,4500.0,9000.0,0.5631111264228821,0.7501111030578613,0.6487501263618469,-0.006703122518956661,0.5692500472068787,0.7573750615119934,0.6548616290092468,0.5139999985694885,0.6920000314712524,0.5998580455780029,,,0.557915985584259,24.163583755493164
-validation,human,1,4500.0,40500.0,0.37967902421951294,0.5846666693687439,0.47633036971092224,-0.07382111996412277,0.3872777819633484,0.5945833325386047,0.48492416739463806,0.31888890266418457,0.5053333640098572,0.40757980942726135,,,0.38891345262527466,27.383806228637695
-validation,human,2,4500.0,36000.0,0.5184444785118103,0.7187777757644653,0.6118195652961731,-0.018171189352869987,0.5287187695503235,0.729593813419342,0.6219512224197388,0.4362500309944153,0.6322500109672546,0.5307666063308716,,,0.516968846321106,27.383806228637695
-validation,human,4,4500.0,27000.0,0.6229629516601562,0.8089629411697388,0.7071906924247742,0.01628497801721096,0.6336249709129333,0.8175833225250244,0.7167206406593323,0.5376666784286499,0.7400000095367432,0.6309514045715332,,,0.6160270571708679,27.383806228637695
-validation,human,5,4500.0,22500.0,0.6493777632713318,0.8219555616378784,0.7280733585357666,0.022814957424998283,0.660099983215332,0.830549955368042,0.7378283143043518,0.5636000037193298,0.7531999945640564,0.6500336527824402,,,0.6369090676307678,27.383806228637695
-validation,human,8,4500.0,9000.0,0.6846666932106018,0.8516666889190674,0.7596060037612915,0.03352390229701996,0.6943750381469727,0.85875004529953,0.7680479884147644,0.6070000529289246,0.7950000166893005,0.6920705437660217,,,0.6696112155914307,27.383806228637695
-test,synthetic,1,500.0,4500.0,0.41688889265060425,0.6495555639266968,0.5274384617805481,-0.07223178446292877,,,,0.41688889265060425,0.6495555639266968,0.5274384617805481,,,0.41443032026290894,23.896039962768555
-test,synthetic,2,500.0,4000.0,0.5692500472068787,0.7830000519752502,0.6672611236572266,0.00026985956355929375,,,,0.5692500472068787,0.7830000519752502,0.6672611236572266,,,0.557541012763977,23.896039962768555
-test,synthetic,4,500.0,3000.0,0.6786666512489319,0.8600000143051147,0.7601169943809509,0.04286628216505051,,,,0.6786666512489319,0.8600000143051147,0.7601169943809509,,,0.6542216539382935,23.896039962768555
-test,synthetic,5,500.0,2500.0,0.7375999689102173,0.9031999707221985,0.8114251494407654,0.06940444558858871,,,,0.7375999689102173,0.9031999707221985,0.8114251494407654,0.748103678226471,0.8272556662559509,0.7113653421401978,23.896039962768555
-test,synthetic,8,500.0,1000.0,0.7070000171661377,0.8790000677108765,0.7862902879714966,0.06344308704137802,,,,0.7070000171661377,0.8790000677108765,0.7862902879714966,,,0.7011929750442505,23.896039962768555
-test,human,1,500.0,4500.0,0.5155555605888367,0.7431111335754395,0.6204044222831726,-0.0215375367552042,,,,0.5155555605888367,0.7431111335754395,0.6204044222831726,,,0.5082485675811768,26.91031265258789
-test,human,2,500.0,4000.0,0.6565000414848328,0.8540000319480896,0.7454702258110046,0.03674145042896271,,,,0.6565000414848328,0.8540000319480896,0.7454702258110046,,,0.6430025100708008,26.91031265258789
-test,human,4,500.0,3000.0,0.737666666507721,0.9136666655540466,0.8140823245048523,0.07469645887613297,,,,0.737666666507721,0.9136666655540466,0.8140823245048523,,,0.726514458656311,26.91031265258789
-test,human,5,500.0,2500.0,0.761199951171875,0.9199999570846558,0.830402135848999,0.08207859843969345,,,,0.761199951171875,0.9199999570846558,0.830402135848999,,,0.7453367114067078,26.91031265258789
-test,human,8,500.0,1000.0,0.7820000648498535,0.9250000715255737,0.8470030426979065,0.08853543549776077,,,,0.7820000648498535,0.9250000715255737,0.8470030426979065,,,0.7604439854621887,26.91031265258789

metrics/test_k_curve.json

@@ -1,204 +0,0 @@
-{
-  "synthetic": {
-    "k1": {
-      "styles": 500.0,
-      "queries": 4500.0,
-      "top1": 0.41688889265060425,
-      "top5": 0.6495555639266968,
-      "mrr": 0.5274384617805481,
-      "margin": -0.07223178446292877,
-      "positive_cosine": 0.5964593291282654,
-      "hardest_negative_cosine": 0.668691098690033,
-      "verification_auroc": 0.41443032026290894,
-      "embedding_std": 0.04344247281551361,
-      "embedding_effective_rank": 23.896039962768555,
-      "unseen_queries": 4500.0,
-      "unseen_top1": 0.41688889265060425,
-      "unseen_top5": 0.6495555639266968,
-      "unseen_mrr": 0.5274384617805481,
-      "unseen_margin": -0.07223178446292877
-    },
-    "k2": {
-      "styles": 500.0,
-      "queries": 4000.0,
-      "top1": 0.5692500472068787,
-      "top5": 0.7830000519752502,
-      "mrr": 0.6672611236572266,
-      "margin": 0.00026985956355929375,
-      "positive_cosine": 0.6695296764373779,
-      "hardest_negative_cosine": 0.6692599058151245,
-      "verification_auroc": 0.557541012763977,
-      "embedding_std": 0.04344247281551361,
-      "embedding_effective_rank": 23.896039962768555,
-      "unseen_queries": 4000.0,
-      "unseen_top1": 0.5692500472068787,
-      "unseen_top5": 0.7830000519752502,
-      "unseen_mrr": 0.6672611236572266,
-      "unseen_margin": 0.00026985956355929375
-    },
-    "k4": {
-      "styles": 500.0,
-      "queries": 3000.0,
-      "top1": 0.6786666512489319,
-      "top5": 0.8600000143051147,
-      "mrr": 0.7601169943809509,
-      "margin": 0.04286628216505051,
-      "positive_cosine": 0.7160320281982422,
-      "hardest_negative_cosine": 0.6731657981872559,
-      "verification_auroc": 0.6542216539382935,
-      "embedding_std": 0.04344247281551361,
-      "embedding_effective_rank": 23.896039962768555,
-      "unseen_queries": 3000.0,
-      "unseen_top1": 0.6786666512489319,
-      "unseen_top5": 0.8600000143051147,
-      "unseen_mrr": 0.7601169943809509,
-      "unseen_margin": 0.04286628216505051
-    },
-    "k5": {
-      "styles": 500.0,
-      "queries": 2500.0,
-      "top1": 0.7375999689102173,
-      "top5": 0.9031999707221985,
-      "mrr": 0.8114251494407654,
-      "margin": 0.06940444558858871,
-      "positive_cosine": 0.7258960008621216,
-      "hardest_negative_cosine": 0.6564915776252747,
-      "verification_auroc": 0.7113653421401978,
-      "embedding_std": 0.04344247281551361,
-      "embedding_effective_rank": 23.896039962768555,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.7375999689102173,
-      "unseen_top5": 0.9031999707221985,
-      "unseen_mrr": 0.8114251494407654,
-      "unseen_margin": 0.06940444558858871,
-      "matched_queries": 500.0,
-      "matched_top1": 0.6600000262260437,
-      "matched_top5": 0.8520000576972961,
-      "matched_mrr": 0.748103678226471,
-      "matched_margin": 0.045967720448970795,
-      "cross_queries": 2000.0,
-      "cross_top1": 0.7570000290870667,
-      "cross_top5": 0.9160000681877136,
-      "cross_mrr": 0.8272556662559509,
-      "cross_margin": 0.0752636268734932
-    },
-    "k8": {
-      "styles": 500.0,
-      "queries": 1000.0,
-      "top1": 0.7070000171661377,
-      "top5": 0.8790000677108765,
-      "mrr": 0.7862902879714966,
-      "margin": 0.06344308704137802,
-      "positive_cosine": 0.739560604095459,
-      "hardest_negative_cosine": 0.676117479801178,
-      "verification_auroc": 0.7011929750442505,
-      "embedding_std": 0.04344247281551361,
-      "embedding_effective_rank": 23.896039962768555,
-      "unseen_queries": 1000.0,
-      "unseen_top1": 0.7070000171661377,
-      "unseen_top5": 0.8790000677108765,
-      "unseen_mrr": 0.7862902879714966,
-      "unseen_margin": 0.06344308704137802
-    }
-  },
-  "human": {
-    "k1": {
-      "styles": 500.0,
-      "queries": 4500.0,
-      "top1": 0.5155555605888367,
-      "top5": 0.7431111335754395,
-      "mrr": 0.6204044222831726,
-      "margin": -0.0215375367552042,
-      "positive_cosine": 0.6388872861862183,
-      "hardest_negative_cosine": 0.6604248285293579,
-      "verification_auroc": 0.5082485675811768,
-      "embedding_std": 0.04376835376024246,
-      "embedding_effective_rank": 26.91031265258789,
-      "unseen_queries": 4500.0,
-      "unseen_top1": 0.5155555605888367,
-      "unseen_top5": 0.7431111335754395,
-      "unseen_mrr": 0.6204044222831726,
-      "unseen_margin": -0.0215375367552042
-    },
-    "k2": {
-      "styles": 500.0,
-      "queries": 4000.0,
-      "top1": 0.6565000414848328,
-      "top5": 0.8540000319480896,
-      "mrr": 0.7454702258110046,
-      "margin": 0.03674145042896271,
-      "positive_cosine": 0.7047365307807922,
-      "hardest_negative_cosine": 0.6679950952529907,
-      "verification_auroc": 0.6430025100708008,
-      "embedding_std": 0.04376835376024246,
-      "embedding_effective_rank": 26.91031265258789,
-      "unseen_queries": 4000.0,
-      "unseen_top1": 0.6565000414848328,
-      "unseen_top5": 0.8540000319480896,
-      "unseen_mrr": 0.7454702258110046,
-      "unseen_margin": 0.03674145042896271
-    },
-    "k4": {
-      "styles": 500.0,
-      "queries": 3000.0,
-      "top1": 0.737666666507721,
-      "top5": 0.9136666655540466,
-      "mrr": 0.8140823245048523,
-      "margin": 0.07469645887613297,
-      "positive_cosine": 0.7484809756278992,
-      "hardest_negative_cosine": 0.6737845540046692,
-      "verification_auroc": 0.726514458656311,
-      "embedding_std": 0.04376835376024246,
-      "embedding_effective_rank": 26.91031265258789,
-      "unseen_queries": 3000.0,
-      "unseen_top1": 0.737666666507721,
-      "unseen_top5": 0.9136666655540466,
-      "unseen_mrr": 0.8140823245048523,
-      "unseen_margin": 0.07469645887613297
-    },
-    "k5": {
-      "styles": 500.0,
-      "queries": 2500.0,
-      "top1": 0.761199951171875,
-      "top5": 0.9199999570846558,
-      "mrr": 0.830402135848999,
-      "margin": 0.08207859843969345,
-      "positive_cosine": 0.7568042874336243,
-      "hardest_negative_cosine": 0.6747256517410278,
-      "verification_auroc": 0.7453367114067078,
-      "embedding_std": 0.04376835376024246,
-      "embedding_effective_rank": 26.91031265258789,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.761199951171875,
-      "unseen_top5": 0.9199999570846558,
-      "unseen_mrr": 0.830402135848999,
-      "unseen_margin": 0.08207859843969345
-    },
-    "k8": {
-      "styles": 500.0,
-      "queries": 1000.0,
-      "top1": 0.7820000648498535,
-      "top5": 0.9250000715255737,
-      "mrr": 0.8470030426979065,
-      "margin": 0.08853543549776077,
-      "positive_cosine": 0.7681199312210083,
-      "hardest_negative_cosine": 0.6795844435691833,
-      "verification_auroc": 0.7604439854621887,
-      "embedding_std": 0.04376835376024246,
-      "embedding_effective_rank": 26.91031265258789,
-      "unseen_queries": 1000.0,
-      "unseen_top1": 0.7820000648498535,
-      "unseen_top5": 0.9250000715255737,
-      "unseen_mrr": 0.8470030426979065,
-      "unseen_margin": 0.08853543549776077
-    }
-  },
-  "checkpoint_score": 0.8209136426448822,
-  "checkpoint_config": {
-    "packed_manifest_sha256": "c1267b9fd45c6e0d222b91d6ae2291b36b763277cc39cbc32484874ffeae9b32",
-    "use_anima": true,
-    "face_dropout": 0.3,
-    "episode_pairs": 8,
-    "seed": 20260719
-  }
-}

metrics/validation_k_curve.json

@@ -1,254 +0,0 @@
-{
-  "synthetic": {
-    "k1": {
-      "styles": 4500.0,
-      "queries": 40500.0,
-      "top1": 0.25738272070884705,
-      "top5": 0.4447901248931885,
-      "mrr": 0.3484785854816437,
-      "margin": -0.13694550096988678,
-      "positive_cosine": 0.6166356205940247,
-      "hardest_negative_cosine": 0.7535811066627502,
-      "verification_auroc": 0.2625024914741516,
-      "embedding_std": 0.04349411278963089,
-      "embedding_effective_rank": 24.163583755493164,
-      "seen_queries": 36000.0,
-      "seen_top1": 0.26072221994400024,
-      "seen_top5": 0.448972225189209,
-      "seen_mrr": 0.35212650895118713,
-      "seen_margin": -0.13477584719657898,
-      "unseen_queries": 4500.0,
-      "unseen_top1": 0.23066666722297668,
-      "unseen_top5": 0.4113333523273468,
-      "unseen_mrr": 0.3192950487136841,
-      "unseen_margin": -0.15430252254009247
-    },
-    "k2": {
-      "styles": 4500.0,
-      "queries": 36000.0,
-      "top1": 0.38269445300102234,
-      "top5": 0.5837222337722778,
-      "mrr": 0.479046106338501,
-      "margin": -0.0735325887799263,
-      "positive_cosine": 0.6810839176177979,
-      "hardest_negative_cosine": 0.7546165585517883,
-      "verification_auroc": 0.3855718970298767,
-      "embedding_std": 0.04349411278963089,
-      "embedding_effective_rank": 24.163583755493164,
-      "seen_queries": 32000.0,
-      "seen_top1": 0.3878125250339508,
-      "seen_top5": 0.5885000228881836,
-      "seen_mrr": 0.48388049006462097,
-      "seen_margin": -0.07155531644821167,
-      "unseen_queries": 4000.0,
-      "unseen_top1": 0.34175002574920654,
-      "unseen_top5": 0.5455000400543213,
-      "unseen_mrr": 0.4403711259365082,
-      "unseen_margin": -0.08935073763132095
-    },
-    "k4": {
-      "styles": 4500.0,
-      "queries": 27000.0,
-      "top1": 0.4913703501224518,
-      "top5": 0.6868888735771179,
-      "mrr": 0.5827372670173645,
-      "margin": -0.032036565244197845,
-      "positive_cosine": 0.7242447733879089,
-      "hardest_negative_cosine": 0.756281316280365,
-      "verification_auroc": 0.4890711009502411,
-      "embedding_std": 0.04349411278963089,
-      "embedding_effective_rank": 24.163583755493164,
-      "seen_queries": 24000.0,
-      "seen_top1": 0.4964166581630707,
-      "seen_top5": 0.6924583315849304,
-      "seen_mrr": 0.5878510475158691,
-      "seen_margin": -0.030110228806734085,
-      "unseen_queries": 3000.0,
-      "unseen_top1": 0.45100000500679016,
-      "unseen_top5": 0.6423333287239075,
-      "unseen_mrr": 0.5418276190757751,
-      "unseen_margin": -0.047447264194488525
-    },
-    "k5": {
-      "styles": 4500.0,
-      "queries": 22500.0,
-      "top1": 0.5625777840614319,
-      "top5": 0.7585333585739136,
-      "mrr": 0.6518997550010681,
-      "margin": -0.006614952813833952,
-      "positive_cosine": 0.7365614771842957,
-      "hardest_negative_cosine": 0.7431764006614685,
-      "verification_auroc": 0.5534452199935913,
-      "embedding_std": 0.04349411278963089,
-      "embedding_effective_rank": 24.163583755493164,
-      "seen_queries": 20000.0,
-      "seen_top1": 0.5681999921798706,
-      "seen_top5": 0.7637499570846558,
-      "seen_mrr": 0.6573799848556519,
-      "seen_margin": -0.004584707319736481,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.5175999999046326,
-      "unseen_top5": 0.7167999744415283,
-      "unseen_mrr": 0.6080579161643982,
-      "unseen_margin": -0.02285691164433956,
-      "matched_queries": 4500.0,
-      "matched_top1": 0.48955556750297546,
-      "matched_top5": 0.6902222037315369,
-      "matched_mrr": 0.5840909481048584,
-      "matched_margin": -0.030851511284708977,
-      "cross_queries": 18000.0,
-      "cross_top1": 0.5808333158493042,
-      "cross_top5": 0.7756111025810242,
-      "cross_mrr": 0.6688519716262817,
-      "cross_margin": -0.0005558129050768912
-    },
-    "k8": {
-      "styles": 4500.0,
-      "queries": 9000.0,
-      "top1": 0.5631111264228821,
-      "top5": 0.7501111030578613,
-      "mrr": 0.6487501263618469,
-      "margin": -0.006703122518956661,
-      "positive_cosine": 0.7508071660995483,
-      "hardest_negative_cosine": 0.7575103044509888,
-      "verification_auroc": 0.557915985584259,
-      "embedding_std": 0.04349411278963089,
-      "embedding_effective_rank": 24.163583755493164,
-      "seen_queries": 8000.0,
-      "seen_top1": 0.5692500472068787,
-      "seen_top5": 0.7573750615119934,
-      "seen_mrr": 0.6548616290092468,
-      "seen_margin": -0.004510539583861828,
-      "unseen_queries": 1000.0,
-      "unseen_top1": 0.5139999985694885,
-      "unseen_top5": 0.6920000314712524,
-      "unseen_mrr": 0.5998580455780029,
-      "unseen_margin": -0.024243783205747604
-    }
-  },
-  "human": {
-    "k1": {
-      "styles": 4500.0,
-      "queries": 40500.0,
-      "top1": 0.37967902421951294,
-      "top5": 0.5846666693687439,
-      "mrr": 0.47633036971092224,
-      "margin": -0.07382111996412277,
-      "positive_cosine": 0.6724859476089478,
-      "hardest_negative_cosine": 0.7463070750236511,
-      "verification_auroc": 0.38891345262527466,
-      "embedding_std": 0.043804656714200974,
-      "embedding_effective_rank": 27.383806228637695,
-      "seen_queries": 36000.0,
-      "seen_top1": 0.3872777819633484,
-      "seen_top5": 0.5945833325386047,
-      "seen_mrr": 0.48492416739463806,
-      "seen_margin": -0.06964229047298431,
-      "unseen_queries": 4500.0,
-      "unseen_top1": 0.31888890266418457,
-      "unseen_top5": 0.5053333640098572,
-      "unseen_mrr": 0.40757980942726135,
-      "unseen_margin": -0.10725177824497223
-    },
-    "k2": {
-      "styles": 4500.0,
-      "queries": 36000.0,
-      "top1": 0.5184444785118103,
-      "top5": 0.7187777757644653,
-      "mrr": 0.6118195652961731,
-      "margin": -0.018171189352869987,
-      "positive_cosine": 0.7339839339256287,
-      "hardest_negative_cosine": 0.7521551847457886,
-      "verification_auroc": 0.516968846321106,
-      "embedding_std": 0.043804656714200974,
-      "embedding_effective_rank": 27.383806228637695,
-      "seen_queries": 32000.0,
-      "seen_top1": 0.5287187695503235,
-      "seen_top5": 0.729593813419342,
-      "seen_mrr": 0.6219512224197388,
-      "seen_margin": -0.014534689486026764,
-      "unseen_queries": 4000.0,
-      "unseen_top1": 0.4362500309944153,
-      "unseen_top5": 0.6322500109672546,
-      "unseen_mrr": 0.5307666063308716,
-      "unseen_margin": -0.04726320877671242
-    },
-    "k4": {
-      "styles": 4500.0,
-      "queries": 27000.0,
-      "top1": 0.6229629516601562,
-      "top5": 0.8089629411697388,
-      "mrr": 0.7071906924247742,
-      "margin": 0.01628497801721096,
-      "positive_cosine": 0.7733137607574463,
-      "hardest_negative_cosine": 0.7570288181304932,
-      "verification_auroc": 0.6160270571708679,
-      "embedding_std": 0.043804656714200974,
-      "embedding_effective_rank": 27.383806228637695,
-      "seen_queries": 24000.0,
-      "seen_top1": 0.6336249709129333,
-      "seen_top5": 0.8175833225250244,
-      "seen_mrr": 0.7167206406593323,
-      "seen_margin": 0.019602350890636444,
-      "unseen_queries": 3000.0,
-      "unseen_top1": 0.5376666784286499,
-      "unseen_top5": 0.7400000095367432,
-      "unseen_mrr": 0.6309514045715332,
-      "unseen_margin": -0.01025400310754776
-    },
-    "k5": {
-      "styles": 4500.0,
-      "queries": 22500.0,
-      "top1": 0.6493777632713318,
-      "top5": 0.8219555616378784,
-      "mrr": 0.7280733585357666,
-      "margin": 0.022814957424998283,
-      "positive_cosine": 0.7811680436134338,
-      "hardest_negative_cosine": 0.7583531141281128,
-      "verification_auroc": 0.6369090676307678,
-      "embedding_std": 0.043804656714200974,
-      "embedding_effective_rank": 27.383806228637695,
-      "seen_queries": 20000.0,
-      "seen_top1": 0.660099983215332,
-      "seen_top5": 0.830549955368042,
-      "seen_mrr": 0.7378283143043518,
-      "seen_margin": 0.026178810745477676,
-      "unseen_queries": 2500.0,
-      "unseen_top1": 0.5636000037193298,
-      "unseen_top5": 0.7531999945640564,
-      "unseen_mrr": 0.6500336527824402,
-      "unseen_margin": -0.004095893818885088
-    },
-    "k8": {
-      "styles": 4500.0,
-      "queries": 9000.0,
-      "top1": 0.6846666932106018,
-      "top5": 0.8516666889190674,
-      "mrr": 0.7596060037612915,
-      "margin": 0.03352390229701996,
-      "positive_cosine": 0.793936014175415,
-      "hardest_negative_cosine": 0.7604120373725891,
-      "verification_auroc": 0.6696112155914307,
-      "embedding_std": 0.043804656714200974,
-      "embedding_effective_rank": 27.383806228637695,
-      "seen_queries": 8000.0,
-      "seen_top1": 0.6943750381469727,
-      "seen_top5": 0.85875004529953,
-      "seen_mrr": 0.7680479884147644,
-      "seen_margin": 0.03656964749097824,
-      "unseen_queries": 1000.0,
-      "unseen_top1": 0.6070000529289246,
-      "unseen_top5": 0.7950000166893005,
-      "unseen_mrr": 0.6920705437660217,
-      "unseen_margin": 0.009157972410321236
-    }
-  },
-  "checkpoint_score": 0.6290457844734192,
-  "checkpoint_config": {
-    "packed_manifest_sha256": "c1267b9fd45c6e0d222b91d6ae2291b36b763277cc39cbc32484874ffeae9b32",
-    "use_anima": true,
-    "face_dropout": 0.3,
-    "episode_pairs": 8,
-    "seed": 20260719
-  }
-}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6cdf126152927f6e31dc856a8d10b8decee3bb20ce4c0ca449206db65d4633ec
-size 69962484
+oid sha256:c7f38c64babe6a1dd711f7819654640b40102d6e73ae45cd2675dc84cda03453
+size 207972492

modeling_anima_style_embedding.py

@@ -1,341 +1,441 @@
 from __future__ import annotations
 
-import math
+from math import isqrt
 
 import torch
-from torch import Tensor, nn
-from torch.nn import functional as F
+from torch import nn
+import torch.nn.functional as F
 
 
-def masked_mean_prototype(per_image: Tensor, mask: Tensor) -> Tensor:
-    if per_image.ndim != 3 or mask.shape != per_image.shape[:2]:
-        raise ValueError("expected per_image [batch,references,width] and matching mask")
-    mask = mask.to(device=per_image.device, dtype=torch.bool)
-    counts = mask.sum(dim=1)
-    if torch.any(counts < 1) or torch.any(counts > 8):
-        raise ValueError("each reference set must contain 1-8 images")
-    weights = mask.unsqueeze(-1).to(per_image.dtype)
-    mean = (per_image * weights).sum(dim=1) / counts.unsqueeze(-1)
-    return F.normalize(mean, dim=-1)
+FACTOR_NAMES = ("line", "color", "texture", "layout")
 
 
-class StyleMixerBlock(nn.Module):
-    def __init__(self, width: int, heads: int, dropout: float = 0.0) -> None:
+class ValueOnlyResampler(nn.Module):
+    """Position-free, bias-free resampling whose output is zero for zero values."""
+
+    def __init__(self, width: int, tokens: int = 144, heads: int = 8) -> None:
         super().__init__()
-        self.query_norm = nn.LayerNorm(width)
-        self.token_norm = nn.LayerNorm(width)
-        self.attention = nn.MultiheadAttention(width, heads, batch_first=True)
-        self.attention_dropout = nn.Dropout(dropout)
-        self.output_norm = nn.LayerNorm(width)
-        self.mlp = nn.Sequential(
-            nn.Linear(width, 4 * width),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(4 * width, width),
-            nn.Dropout(dropout),
-        )
+        if width % heads:
+            raise ValueError("width must be divisible by heads")
+        self.heads = heads
+        self.head_dim = width // heads
+        self.queries = nn.Parameter(torch.empty(tokens, width))
+        self.q_proj = nn.Linear(width, width, bias=False)
+        self.k_proj = nn.Linear(width, width, bias=False)
+        self.v_proj = nn.Linear(width, width, bias=False)
+        self.out_proj = nn.Linear(width, width, bias=False)
+        nn.init.normal_(self.queries, std=width**-0.5)
 
-    def forward(self, query: Tensor, tokens: Tensor, token_mask: Tensor) -> Tensor:
-        value, _ = self.attention(
-            self.query_norm(query),
-            self.token_norm(tokens),
-            self.token_norm(tokens),
-            key_padding_mask=~token_mask,
-            need_weights=False,
-        )
-        query = query + self.attention_dropout(value)
-        return query + self.mlp(self.output_norm(query))
+    def forward(self, values: torch.Tensor) -> torch.Tensor:
+        batch, source_tokens, width = values.shape
+        target_tokens = self.queries.shape[0]
+        queries = self.queries.unsqueeze(0).expand(batch, -1, -1)
+        q = self.q_proj(queries).view(batch, target_tokens, self.heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(values).view(batch, source_tokens, self.heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(values).view(batch, source_tokens, self.heads, self.head_dim).transpose(1, 2)
+        result = F.scaled_dot_product_attention(q, k, v)
+        return self.out_proj(result.transpose(1, 2).reshape(batch, target_tokens, width))
 
 
-class UnifiedStyleEncoder(nn.Module):
-    """Encode cached full/face vision tokens and optional Anima descriptors."""
+class SharedAxisStyleEncoder(nn.Module):
+    """One VAE field and one shared axis space for embedding, routing, and transfer."""
 
     def __init__(
         self,
-        *,
-        external_dim: int = 1152,
-        external_tokens: int = 30,
-        anima_dim: int = 4096,
-        anima_tokens: int = 3,
-        width: int = 384,
-        blocks: int = 2,
-        heads: int = 6,
-        embedding_dim: int = 512,
-        use_anima: bool = True,
-        dropout: float = 0.0,
-        reference_grounding: bool = False,
-        functional_factors: bool = False,
+        dino_dim: int,
+        style_dim: int = 1024,
+        axis_count: int = 256,
+        embedding_dim: int = 1024,
+        anima_dim: int = 0,
+        resolution: str = "s12",
+        set_layers: int = 2,
+        set_heads: int = 8,
+        direct_context_axes: bool = False,
+        axis_aligned: bool = False,
+        disable_token_film: bool = False,
+        translation_invariant_stats: bool = False,
     ) -> None:
         super().__init__()
-        if (
-            blocks < 1
-            or width % heads
-            or (reference_grounding and embedding_dim % heads)
-            or not 0.0 <= dropout < 1.0
-        ):
-            raise ValueError("blocks must be positive and width must be divisible by heads")
-        if functional_factors and embedding_dim != 512:
-            raise ValueError("functional factorization currently requires a 512-D embedding")
-        self.external_tokens = external_tokens
-        self.anima_tokens = anima_tokens if use_anima else 0
-        self.external_projector = nn.Sequential(
-            nn.LayerNorm(external_dim),
-            nn.Linear(external_dim, width),
+        if resolution not in {"s12", "s24r"}:
+            raise ValueError(f"unknown spatial resolution: {resolution}")
+        if axis_count % set_heads:
+            raise ValueError("axis_count must be divisible by set_heads")
+        self.resolution = resolution
+        self.style_dim = style_dim
+        self.axis_count = axis_count
+        self.embedding_dim = embedding_dim
+        self.set_layers = set_layers
+        self.set_heads = set_heads
+        self.direct_context_axes = direct_context_axes
+        self.axis_aligned = axis_aligned
+        self.disable_token_film = disable_token_film
+        self.translation_invariant_stats = translation_invariant_stats
+        if disable_token_film and not axis_aligned:
+            raise ValueError("disable_token_film requires axis_aligned=True")
+        if translation_invariant_stats and not axis_aligned:
+            raise ValueError("translation-invariant statistics require axis_aligned=True")
+        last_stride = 2 if resolution == "s12" else 1
+        self.stem = nn.Sequential(
+            nn.Conv2d(16, 256, 3, stride=2, padding=1),
+            nn.GELU(),
+            nn.Conv2d(256, 512, 3, stride=2, padding=1),
+            nn.GELU(),
+            nn.Conv2d(512, style_dim, 3, stride=last_stride, padding=1),
         )
-        self.anima_projector = (
-            nn.Sequential(nn.LayerNorm(anima_dim), nn.Linear(anima_dim, width))
-            if use_anima
-            else None
+        self.token_norm = nn.LayerNorm(style_dim, elementwise_affine=False)
+        self.dino_norm = nn.LayerNorm(dino_dim)
+        self.dino_film = nn.Linear(dino_dim, 2 * style_dim, bias=False)
+        self.dino_relation = nn.Linear(dino_dim, axis_count, bias=False)
+        nn.init.zeros_(self.dino_film.weight)
+
+        self.axis_dictionary = nn.Parameter(torch.empty(style_dim, axis_count))
+        nn.init.orthogonal_(self.axis_dictionary)
+        self.membership_logits = nn.Parameter(torch.zeros(len(FACTOR_NAMES), axis_count))
+        self.moment_projection = nn.Linear(2 * axis_count, axis_count, bias=False)
+        relation_layer = nn.TransformerEncoderLayer(
+            axis_count,
+            nhead=set_heads,
+            dim_feedforward=2 * axis_count,
+            dropout=0.0,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,
         )
-        self.full_positions = nn.Parameter(torch.randn(external_tokens, width) * 0.02)
-        self.face_positions = nn.Parameter(torch.randn(external_tokens, width) * 0.02)
-        self.anima_positions = (
-            nn.Parameter(torch.randn(anima_tokens, width) * 0.02) if use_anima else None
+        self.relation_encoder = nn.TransformerEncoder(
+            relation_layer,
+            set_layers,
+            enable_nested_tensor=False,
         )
-        self.reference_grounding = reference_grounding
-        self.functional_factors = functional_factors
-        query_count = 5 if functional_factors else 1
-        self.style_query = nn.Parameter(torch.randn(1, query_count, width) * 0.02)
-        self.blocks = nn.ModuleList(
-            [StyleMixerBlock(width, heads, dropout) for _ in range(blocks)]
+        self.axis_reliability = nn.Linear(axis_count, axis_count)
+        self.face_reliability = nn.Linear(axis_count, 1)
+        self.embedding_head = nn.Sequential(
+            nn.LayerNorm(axis_count),
+            nn.Linear(axis_count, embedding_dim),
         )
-        if functional_factors:
-            self.factor_output = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, 96))
-            self.shared_output = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, 128))
-            self.output = nn.Sequential(nn.LayerNorm(512), nn.Linear(512, embedding_dim))
-            self.factor_decoders = nn.ModuleList([nn.Linear(96, external_dim) for _ in range(4)])
-        else:
-            self.output = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, embedding_dim))
-        if reference_grounding:
-            self.consensus_query = nn.Parameter(torch.randn(1, 1, embedding_dim) * 0.02)
-            self.consensus_attention = nn.MultiheadAttention(
-                embedding_dim, heads, batch_first=True, dropout=dropout
-            )
-            self.consensus_condition = nn.Linear(embedding_dim, width)
-            self.consensus_gamma = nn.Parameter(torch.tensor(-2.9444))
-
-    def _read_tokens(
-        self, tokens: Tensor, token_mask: Tensor, conditioning: Tensor | None = None
-    ) -> tuple[Tensor, Tensor | None]:
-        query = self.style_query.expand(tokens.shape[0], -1, -1)
-        if conditioning is not None:
-            query = query + self.consensus_condition(conditioning)[:, None]
-        for block in self.blocks:
-            query = block(query, tokens, token_mask)
-        if self.functional_factors:
-            factors = F.normalize(self.factor_output(query[:, :4]), dim=-1)
-            shared = self.shared_output(query[:, 4])
-            embedding = self.output(torch.cat((factors.flatten(1), shared), dim=-1))
-            return F.normalize(embedding, dim=-1), factors
-        return F.normalize(self.output(query[:, 0]), dim=-1), None
-
-    def encode_images(
-        self,
-        full_features: Tensor,
-        *,
-        face_features: Tensor | None = None,
-        face_mask: Tensor | None = None,
-        anima_features: Tensor | None = None,
-        conditioning: Tensor | None = None,
-    ) -> Tensor:
-        """Encode independent images without padding them into reference sets."""
-
-        if full_features.ndim != 3 or full_features.shape[1] != self.external_tokens:
-            raise ValueError("full_features must have shape [images,external_tokens,width]")
-        images = full_features.shape[0]
-        full = self.external_projector(full_features) + self.full_positions
-        tokens = [full]
-        masks = [torch.ones(full.shape[:2], dtype=torch.bool, device=full.device)]
-
-        if face_features is not None:
-            if face_features.shape != full_features.shape:
-                raise ValueError("face_features must match full_features")
-            if face_mask is None or face_mask.shape != (images,):
-                raise ValueError("face_mask is required and must contain one value per image")
-            face = self.external_projector(face_features)
-            tokens.append(face + self.face_positions)
-            masks.append(face_mask[:, None].to(torch.bool).expand(-1, self.external_tokens))
-
-        if self.anima_projector is not None:
-            expected = (images, self.anima_tokens)
-            if anima_features is None or anima_features.shape[:2] != expected:
-                raise ValueError("anima_features are required by this encoder configuration")
-            anima = self.anima_projector(anima_features)
-            tokens.append(anima + self.anima_positions)
-            masks.append(torch.ones(anima.shape[:2], dtype=torch.bool, device=anima.device))
-
-        token_tensor = torch.cat(tokens, dim=1)
-        token_mask = torch.cat(masks, dim=1)
-        return self._read_tokens(token_tensor, token_mask, conditioning)[0]
-
-    def encode_images_with_factors(
-        self,
-        full_features: Tensor,
-        *,
-        face_features: Tensor | None = None,
-        face_mask: Tensor | None = None,
-        anima_features: Tensor | None = None,
-        conditioning: Tensor | None = None,
-    ) -> tuple[Tensor, Tensor | None]:
-        if full_features.ndim != 3 or full_features.shape[1] != self.external_tokens:
-            raise ValueError("full_features must have shape [images,external_tokens,width]")
-        images = full_features.shape[0]
-        full = self.external_projector(full_features) + self.full_positions
-        tokens = [full]
-        masks = [torch.ones(full.shape[:2], dtype=torch.bool, device=full.device)]
-        if face_features is not None:
-            if face_features.shape != full_features.shape:
-                raise ValueError("face_features must match full_features")
-            if face_mask is None or face_mask.shape != (images,):
-                raise ValueError("face_mask is required and must contain one value per image")
-            tokens.append(self.external_projector(face_features) + self.face_positions)
-            masks.append(face_mask[:, None].to(torch.bool).expand(-1, self.external_tokens))
-        if self.anima_projector is not None:
-            if anima_features is None or anima_features.shape[:2] != (images, self.anima_tokens):
-                raise ValueError("anima_features are required by this encoder configuration")
-            tokens.append(self.anima_projector(anima_features) + self.anima_positions)
-            masks.append(torch.ones(anima_features.shape[:2], dtype=torch.bool, device=full.device))
-        return self._read_tokens(torch.cat(tokens, dim=1), torch.cat(masks, dim=1), conditioning)
-
-    def encode_reference_sets(
-        self,
-        full_features: Tensor,
-        reference_mask: Tensor,
-        *,
-        face_features: Tensor | None = None,
-        face_mask: Tensor | None = None,
-        anima_features: Tensor | None = None,
-    ) -> tuple[Tensor, Tensor, Tensor | None]:
-        batch, references = full_features.shape[:2]
-        flat_face = None if face_features is None else face_features.flatten(0, 1)
-        flat_face_mask = None if face_mask is None else face_mask.flatten()
-        flat_anima = None if anima_features is None else anima_features.flatten(0, 1)
-        per_image, factors = self.encode_images_with_factors(
-            full_features.flatten(0, 1),
-            face_features=flat_face,
-            face_mask=flat_face_mask,
-            anima_features=flat_anima,
+        self.resampler = (
+            ValueOnlyResampler(style_dim, tokens=144, heads=set_heads)
+            if resolution == "s24r"
+            else None
         )
-        per_image = per_image.reshape(batch, references, -1)
-        factors = None if factors is None else factors.reshape(batch, references, 4, -1)
-        base = masked_mean_prototype(per_image, reference_mask)
-        if not self.reference_grounding:
-            return base, per_image, factors
-        consensus, _ = self.consensus_attention(
-            self.consensus_query.expand(batch, -1, -1),
-            per_image,
-            per_image,
-            key_padding_mask=~reference_mask,
-            need_weights=False,
+        # Construct optional modality modules last so matched runs share identical common initialization.
+        self.anima_norm = nn.LayerNorm(anima_dim) if anima_dim else None
+        self.anima_film = nn.Linear(anima_dim, 2 * style_dim, bias=False) if anima_dim else None
+        self.anima_relation = nn.Linear(anima_dim, axis_count, bias=False) if anima_dim else None
+        if self.anima_film is not None:
+            nn.init.zeros_(self.anima_film.weight)
+        if axis_aligned:
+            self.axis_moment_weights = nn.Parameter(torch.empty(axis_count, 2))
+            nn.init.normal_(self.axis_moment_weights, std=axis_count**-0.5)
+        if translation_invariant_stats:
+            self.axis_power_weights = nn.Parameter(torch.zeros(axis_count, 4))
+
+    def model_config(self) -> dict[str, int | str | bool]:
+        return {
+            "style_dim": self.style_dim,
+            "axis_count": self.axis_count,
+            "embedding_dim": self.embedding_dim,
+            "anima_dim": self.anima_norm.normalized_shape[0] if self.anima_norm is not None else 0,
+            "resolution": self.resolution,
+            "set_layers": self.set_layers,
+            "set_heads": self.set_heads,
+            "direct_context_axes": self.direct_context_axes,
+            "axis_aligned": self.axis_aligned,
+            "disable_token_film": self.disable_token_film,
+            "translation_invariant_stats": self.translation_invariant_stats,
+        }
+
+    def _field(
+        self,
+        latent: torch.Tensor,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
+        axis_shift: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch, references = latent.shape[:2]
+        field = self.stem(latent.flatten(0, 1).float())
+        field = field.flatten(2).transpose(1, 2).reshape(batch, references, -1, field.shape[1])
+        field = self.token_norm(field)
+        field = field * (1 + scale.unsqueeze(2)) + shift.unsqueeze(2)
+        axes = torch.einsum("brnd,dk->brnk", field, self.axis_dictionary.to(field.dtype))
+        if axis_shift is not None:
+            axes = axes + axis_shift.unsqueeze(2)
+        return field, axes
+
+    def _moments(self, axes: torch.Tensor) -> torch.Tensor:
+        mean = axes.mean(dim=2)
+        log_std = torch.log(axes.std(dim=2, unbiased=False).clamp_min(1e-6))
+        if self.axis_aligned:
+            coordinates = (
+                torch.stack((mean, log_std), dim=-1)
+                * self.axis_moment_weights.to(mean.dtype)
+            ).sum(-1)
+            if self.translation_invariant_stats:
+                power = self._power_statistics(axes).to(mean.dtype)
+                coordinates = coordinates + (
+                    power * self.axis_power_weights.to(mean.dtype)
+                ).sum(-1)
+            return coordinates
+        moments = torch.cat((mean, log_std), dim=-1)
+        return self.moment_projection(moments)
+
+    def _power_statistics(self, axes: torch.Tensor) -> torch.Tensor:
+        """Phase-free radial and directional power for each aligned spatial axis."""
+        side = isqrt(axes.shape[2])
+        if side * side != axes.shape[2]:
+            raise ValueError(f"unexpected spatial token count: {axes.shape[2]}")
+        fy = torch.fft.fftfreq(side, device=axes.device)[:, None]
+        fx = torch.fft.rfftfreq(side, device=axes.device)[None, :]
+        radius2 = fy.square() + fx.square()
+        radius = radius2.sqrt()
+        basis = torch.stack((
+            ((radius > 0) & (radius <= 0.25)).float(),
+            ((radius > 0.25) & (radius <= 0.5)).float(),
+            (fx.square() - fy.square()) / radius2.clamp_min(1e-12),
+            4 * fx.square() * fy.square() / radius2.square().clamp_min(1e-12),
+        ))
+        rfft_weight = torch.ones(side // 2 + 1, device=axes.device)
+        rfft_weight[1:-1] = 2
+        maps = axes.float().permute(0, 1, 3, 2).reshape(
+            *axes.shape[:2],
+            axes.shape[-1],
+            side,
+            side,
         )
-        consensus = F.normalize(consensus[:, 0], dim=-1)
-        reread, factors = self.encode_images_with_factors(
-            full_features.flatten(0, 1),
-            face_features=flat_face,
-            face_mask=flat_face_mask,
-            anima_features=flat_anima,
-            conditioning=consensus[:, None].expand(-1, references, -1).reshape(batch * references, -1),
+        maps = maps - maps.mean(dim=(-2, -1), keepdim=True)
+        power = torch.fft.rfft2(maps, norm="ortho").abs().square()
+        power = power * rfft_weight[None, None, None, None, :]
+        power = power / power.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-12)
+        return torch.einsum("brkhw,phw->brkp", power, basis)
+
+    def factor_coordinates(self, axis_coordinates: torch.Tensor) -> torch.Tensor:
+        """Apply overlapping factor memberships without creating factor branches."""
+        memberships = torch.sigmoid(self.membership_logits).to(axis_coordinates.dtype)
+        return axis_coordinates.unsqueeze(-2) * memberships
+
+    def encode_vae_axes(self, latents: torch.Tensor) -> torch.Tensor:
+        """Encode VAE latents without DINO/Anima context for the frozen transfer teacher."""
+        if latents.ndim != 5 or latents.shape[2] != 16:
+            raise ValueError(f"expected latents [B,R,16,H,W], got {tuple(latents.shape)}")
+        batch, references = latents.shape[:2]
+        zeros = torch.zeros(
+            batch,
+            references,
+            self.axis_dictionary.shape[0],
+            device=latents.device,
+            dtype=latents.dtype,
         )
-        reread = reread.reshape(batch, references, -1)
-        factors = None if factors is None else factors.reshape(batch, references, 4, -1)
-        agreement = F.cosine_similarity(reread, consensus[:, None], dim=-1)
-        centered = agreement - (
-            (agreement * reference_mask).sum(1, keepdim=True)
-            / reference_mask.sum(1, keepdim=True).clamp_min(1)
+        _, axes = self._field(latents, zeros, zeros)
+        return self._moments(axes)
+
+    def _user_axis_gate(self, user_weights: torch.Tensor) -> torch.Tensor:
+        memberships = torch.sigmoid(self.membership_logits)
+        factors = user_weights[..., : len(FACTOR_NAMES)].unsqueeze(-1)
+        return 1 - torch.prod(1 - factors * memberships[None, None], dim=-2)
+
+    def _controlled_tokens(
+        self,
+        axes: torch.Tensor,
+        beta: torch.Tensor,
+        view_gate: torch.Tensor,
+    ) -> torch.Tensor:
+        gated = axes * beta.unsqueeze(2) * view_gate[:, :, None, None]
+        tokens = torch.einsum(
+            "brnk,dk->brnd",
+            gated,
+            self.axis_dictionary.to(gated.dtype),
         )
-        weights = (1.0 + 0.25 * torch.tanh(4.0 * centered)) * reference_mask
-        grounded = F.normalize((reread * weights[..., None]).sum(1) / weights.sum(1, keepdim=True), dim=-1)
-        gamma = torch.sigmoid(self.consensus_gamma)
-        return F.normalize(base + gamma * (grounded - base), dim=-1), reread, factors
-
-    def functional_response_loss(self, factor_codes: Tensor, full_features: Tensor) -> Tensor:
-        """Reconstruct artist response after removing a shared prompt/seed cell mean."""
-        if not self.functional_factors:
-            return full_features.new_zeros(())
-        groups = ((2, 3, 4, 5), (0, 6, 24), tuple(range(8, 24)), (1, 7, 25))
-        losses = []
-        for index, token_indices in enumerate(groups):
-            teacher = full_features[..., list(token_indices), :].float().mean(dim=-2)
-            teacher = F.normalize(teacher - teacher.mean(dim=1, keepdim=True), dim=-1)
-            prediction = F.normalize(self.factor_decoders[index](factor_codes[..., index, :]).float(), dim=-1)
-            losses.append((1.0 - F.cosine_similarity(prediction, teacher, dim=-1)).mean())
-        return torch.stack(losses).mean()
+        if self.resampler is not None:
+            batch, references, token_count, width = tokens.shape
+            tokens = self.resampler(tokens.reshape(batch * references, token_count, width))
+            tokens = tokens.reshape(batch, references, -1, width)
+        return tokens
+
+    def soft_orthogonality_loss(self) -> torch.Tensor:
+        dictionary = F.normalize(self.axis_dictionary, dim=0)
+        gram = dictionary.T @ dictionary
+        return (gram - torch.eye(gram.shape[0], device=gram.device, dtype=gram.dtype)).square().mean()
 
     def forward(
         self,
-        full_features: Tensor,
-        reference_mask: Tensor,
-        *,
-        face_features: Tensor | None = None,
-        face_mask: Tensor | None = None,
-        anima_features: Tensor | None = None,
-    ) -> tuple[Tensor, Tensor]:
-        if full_features.ndim != 4 or full_features.shape[2] != self.external_tokens:
-            raise ValueError("full_features must have shape [batch,references,external_tokens,width]")
-        batch, references = full_features.shape[:2]
-        if reference_mask.shape != (batch, references):
-            raise ValueError("reference_mask shape does not match full_features")
-
-        prototype, per_image, _ = self.encode_reference_sets(
-            full_features,
-            reference_mask,
-            face_features=face_features,
-            face_mask=face_mask,
-            anima_features=anima_features,
+        full_latents: torch.Tensor,
+        dino: torch.Tensor,
+        reference_valid: torch.Tensor,
+        face_latents: torch.Tensor | None = None,
+        face_valid: torch.Tensor | None = None,
+        user_weights: torch.Tensor | None = None,
+        mode: str = "auto",
+        overall_style_gain: float | torch.Tensor = 1.0,
+        build_memory: bool = True,
+        anima: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        if full_latents.ndim != 5 or full_latents.shape[2] != 16:
+            raise ValueError(f"expected full latents [B,R,16,H,W], got {tuple(full_latents.shape)}")
+        if not reference_valid.any(dim=1).all():
+            raise ValueError("every sample needs at least one valid reference")
+        if mode not in {"auto", "assisted", "manual"}:
+            raise ValueError(f"unknown routing mode: {mode}")
+
+        batch, references = full_latents.shape[:2]
+        if user_weights is None:
+            user_weights = torch.ones(
+                batch,
+                references,
+                len(FACTOR_NAMES) + 1,
+                device=full_latents.device,
+                dtype=full_latents.dtype,
+            )
+        context = self.dino_norm(dino.float())
+        scale, shift = self.dino_film(context).chunk(2, dim=-1)
+        relation_context = self.dino_relation(context)
+        if self.anima_norm is not None:
+            if anima is None:
+                raise ValueError("Anima features are required when anima_dim is configured")
+            anima_context = self.anima_norm(anima.float().flatten(start_dim=2))
+            anima_scale, anima_shift = self.anima_film(anima_context).chunk(2, dim=-1)
+            scale = scale + anima_scale
+            shift = shift + anima_shift
+            relation_context = relation_context + self.anima_relation(anima_context)
+        elif anima is not None:
+            raise ValueError("Anima features were provided but anima_dim=0")
+        if self.disable_token_film:
+            scale = torch.zeros_like(scale)
+            shift = torch.zeros_like(shift)
+        axis_shift = (
+            relation_context
+            if self.axis_aligned and self.direct_context_axes
+            else None
         )
-        return prototype, per_image
+        full_field, full_axes = self._field(
+            full_latents,
+            scale,
+            shift,
+            axis_shift,
+        )
+        full_coordinates = self._moments(full_axes)
+        if self.direct_context_axes and not self.axis_aligned:
+            full_coordinates = full_coordinates + relation_context
 
+        if face_latents is None:
+            face_valid = torch.zeros_like(reference_valid)
+            face_axes = None
+            face_coordinates = torch.zeros_like(full_coordinates)
+        else:
+            if face_valid is None:
+                raise ValueError("face_valid is required with face_latents")
+            _, face_axes = self._field(
+                face_latents,
+                scale,
+                shift,
+                axis_shift,
+            )
+            face_coordinates = self._moments(face_axes)
 
-class AngularPrototypeLoss(nn.Module):
-    """Episode prototype softmax with an optional ArcFace-style positive margin."""
+        preliminary = (
+            full_coordinates
+            + face_coordinates * face_valid.unsqueeze(-1)
+            + (0 if self.direct_context_axes else relation_context)
+        )
+        contextual = self.relation_encoder(preliminary, src_key_padding_mask=~reference_valid)
+        reliability = torch.sigmoid(self.axis_reliability(contextual))
+        predicted_face = torch.sigmoid(self.face_reliability(contextual).squeeze(-1))
 
-    def __init__(
-        self, initial_scale: float = 10.0, initial_bias: float | None = None
-    ) -> None:
-        super().__init__()
-        if initial_scale <= 0:
-            raise ValueError("initial_scale must be positive")
-        self.raw_scale = nn.Parameter(torch.tensor(math.log(math.expm1(initial_scale))))
-        # initial_bias remains accepted for feasibility-script compatibility. A
-        # common softmax bias is a no-op and is not a trainable production parameter.
+        if mode == "auto":
+            face_weight = predicted_face
+        elif mode == "assisted":
+            face_weight = predicted_face * user_weights[..., -1]
+        else:
+            face_weight = user_weights[..., -1]
+        face_weight = face_weight * face_valid * reference_valid
+        coordinates = (
+            full_coordinates + face_weight.unsqueeze(-1) * face_coordinates
+        ) / (1 + face_weight.unsqueeze(-1))
 
-    def forward(
-        self,
-        queries: Tensor,
-        prototypes: Tensor,
-        targets: Tensor,
-        *,
-        margin: float = 0.0,
-    ) -> tuple[Tensor, Tensor]:
-        if queries.ndim < 2 or prototypes.ndim != queries.ndim:
-            raise ValueError("queries and prototypes must have matching [episode,...,width] ranks")
-        if queries.shape[:-2] != prototypes.shape[:-2] or queries.shape[-1] != prototypes.shape[-1]:
-            raise ValueError("queries and prototypes must have matching episode axes and widths")
-        if targets.shape != queries.shape[:-1]:
-            raise ValueError("targets must contain one class index per query")
-        # Autocast can round a near-perfect BF16 cosine to exactly one. ArcFace's
-        # sqrt(1-cos��) derivative is then singular, so margin math and CE stay FP32.
-        cosine = (
-            F.normalize(queries, dim=-1)
-            @ F.normalize(prototypes, dim=-1).transpose(-1, -2)
-        ).float()
-        loss_cosine = cosine
-        if margin:
-            if not 0.0 <= margin < math.pi / 2:
-                raise ValueError("margin must be in [0, pi/2)")
-            positive = cosine.gather(-1, targets.unsqueeze(-1)).clamp(-1 + 1e-6, 1 - 1e-6)
-            positive_margin = positive * math.cos(margin) - torch.sqrt(1 - positive.square()) * math.sin(margin)
-            loss_cosine = cosine.scatter(-1, targets.unsqueeze(-1), positive_margin)
-        scale = F.softplus(self.raw_scale).clamp_max(100.0)
-        loss_logits = scale * loss_cosine
-        metric_logits = scale * cosine
-        return (
-            F.cross_entropy(loss_logits.flatten(0, -2), targets.flatten()),
-            metric_logits,
+        axis_gate = self._user_axis_gate(user_weights.float())
+        valid = reference_valid.unsqueeze(-1).to(axis_gate.dtype)
+        if mode == "auto":
+            weights = reliability * valid
+            gamma = torch.ones(batch, self.axis_dictionary.shape[1], device=weights.device, dtype=weights.dtype)
+        elif mode == "assisted":
+            weights = reliability * axis_gate * valid
+            gamma = axis_gate.amax(dim=1)
+        else:
+            weights = axis_gate * valid
+            gamma = axis_gate.amax(dim=1)
+
+        denominator = weights.sum(dim=1, keepdim=True)
+        alpha = torch.where(
+            denominator > 0,
+            weights / denominator.clamp_min(torch.finfo(weights.dtype).tiny),
+            torch.zeros_like(weights),
         )
+        beta = alpha * gamma.unsqueeze(1)
+        mixed = gamma * torch.einsum("brk,brk->bk", alpha, coordinates)
+
+        gain = torch.as_tensor(overall_style_gain, device=mixed.device, dtype=mixed.dtype)
+        while gain.ndim < mixed.ndim:
+            gain = gain.unsqueeze(-1)
+        mixed = mixed * gain
+        beta = beta * gain.unsqueeze(1) if gain.ndim == 2 else beta * gain
+
+        pre_embedding = self.embedding_head(mixed)
+        output = {
+            "embedding": F.normalize(pre_embedding, dim=-1),
+            "pre_embedding": pre_embedding,
+            "style_present": mixed.ne(0).any(dim=-1),
+            "reference_axis_coordinates": coordinates,
+            "axis_coordinates": mixed,
+            "axis_reliability": reliability,
+            "axis_weights": weights,
+            "axis_alpha": alpha,
+            "axis_gamma": gamma,
+            "factor_memberships": torch.sigmoid(self.membership_logits),
+            "reference_factor_coordinates": self.factor_coordinates(coordinates),
+            "factor_coordinates": self.factor_coordinates(mixed),
+            "face_weight": face_weight,
+        }
+        if build_memory:
+            full_tokens = self._controlled_tokens(
+                full_axes,
+                beta,
+                reference_valid.to(beta.dtype),
+            )
+            memories = [full_tokens.flatten(1, 2)]
+            memory_masks = [
+                reference_valid[:, :, None].expand(-1, -1, full_tokens.shape[2]).flatten(1)
+            ]
+            if face_axes is not None:
+                face_tokens = self._controlled_tokens(face_axes, beta, face_weight)
+                memories.append(face_tokens.flatten(1, 2))
+                memory_masks.append(
+                    (reference_valid & face_valid)[:, :, None]
+                    .expand(-1, -1, face_tokens.shape[2])
+                    .flatten(1)
+                )
+            output.update({
+                "style_memory": torch.cat(memories, dim=1),
+                "style_memory_valid": torch.cat(memory_masks, dim=1),
+                "style_condition": torch.einsum(
+                    "bk,dk->bd",
+                    mixed,
+                    self.axis_dictionary.to(mixed.dtype),
+                ),
+            })
+        return output
 
 
-def subset_consistency_loss(first: Tensor, second: Tensor) -> Tensor:
-    if first.shape != second.shape or first.ndim != 2:
-        raise ValueError("subset prototypes must have equal [batch,width] shapes")
-    return (1.0 - F.cosine_similarity(first, second, dim=-1)).mean()
+def load_shared_axis_state(
+    model: SharedAxisStyleEncoder,
+    state: dict[str, torch.Tensor],
+) -> None:
+    """Load a checkpoint, deriving aligned per-axis moments from a legacy dense head."""
+    state = dict(state)
+    if model.axis_aligned and "axis_moment_weights" not in state:
+        dense = state["moment_projection.weight"]
+        axes = model.axis_count
+        state["axis_moment_weights"] = torch.stack((
+            dense.diagonal(),
+            dense[:, axes:].diagonal(),
+        ), dim=-1)
+    if model.translation_invariant_stats and "axis_power_weights" not in state:
+        state["axis_power_weights"] = torch.zeros_like(model.axis_power_weights)
+    model.load_state_dict(state)

requirements.txt

@@ -1,3 +1,3 @@
-torch>=2.4
-safetensors>=0.4
-huggingface_hub>=0.28
+torch>=2.5
+safetensors>=0.4.5
+huggingface_hub>=0.27