Explain the winning architecture, exact distillation loss, and inference path

Document active and inactive losses, frozen versus updated parameters, mask calibration, and the expanded-source alternative. Refresh documentation checksums; all checkpoint weights are unchanged.

FORMULATION.md

@@ -0,0 +1,213 @@
+# The released SD1.5 → SANA winner: architecture, loss, and inference
+
+The default [released checkpoint](https://huggingface.co/ntc-ai/model-glue-sd15-sana-text)
+is a **7,182,917-parameter text-conditioning bridge**. Its winning continuation
+trained the final linear output projection using paired embeddings and frozen-SANA
+response distillation. The parent came from earlier ParticleGAN experiments,
+but this continuation has **no active GAN or b_cap loss**. Its learned particle
+cloud is retained as fixed network weights.
+
+`model.safetensors` is the EMA checkpoint selected at update 1,250 of 1,500,
+followed by one validation-fitted mask-bias adjustment. The uncalibrated control
+and expanded-source alternative are also released. See the [results](https://github.com/255BITS/model-glue/blob/main/docs/results/formulation-retraining-20260905.md),
+[download instructions](https://github.com/255BITS/model-glue/blob/main/docs/sana-transfer.md#released-sd15-text-checkpoints), and
+[exact recipe](https://github.com/255BITS/model-glue/blob/main/configs/formulation/linear-distill-fixed.json).
+
+## Architecture
+
+```mermaid
+flowchart LR
+    A[Prompt] --> B[Frozen SD1.5 CLIP]
+    B --> C[77 x 768 contextual states]
+    C --> D[LayerNorm and 768-to-384 projection]
+    C --> R[Mean pooling and particle router]
+    P[Fixed 256 x 4 learned cloud] --> R
+    D --> Q[77 learned queries plus source and particle offsets]
+    R --> Q
+    Q --> T[4 cross-attention blocks; width 384]
+    D --> T
+    T --> O[Linear output and fixed feature calibration]
+    T --> M[Linear mask head and threshold]
+    O --> S[Frozen SANA: 20 sampling steps]
+    M --> S
+    U[Stored native empty conditioning] --> S
+    S --> V[Frozen DC-AE decoder]
+    V --> I[Image]
+```
+
+The input is the full contextual CLIP sequence, not a pooled embedding. For
+source states $X$ with shape 77 × 768:
+
+$$
+E=\operatorname{Linear}_{768\to384}(\operatorname{LN}(X)),\qquad
+r=\operatorname{MLP}_{768\to64\to4}(\operatorname{mean}_{tokens}X).
+$$
+
+The router MLP uses SiLU. With saved cloud $P$ of shape 256 × 4:
+
+$$
+a=\operatorname{softmax}(rP^T/\sqrt4),\qquad p=aP,\qquad
+H_0=Q+E+\operatorname{Linear}_{4\to384}(p).
+$$
+
+$Q$ contains 77 learned queries of width 384. The same four-dimensional $p$
+provides a global query offset. This is attention over shared parameter vectors,
+not a lookup of training captions. The cloud and router stayed fixed during the
+winning continuation; no particle advantage is established.
+
+Each of four blocks performs six-head cross-attention to $E$, then a residual
+feed-forward network:
+
+$$
+A_l=H_l+\operatorname{MHA}(\operatorname{LN}(H_l),E,E),\qquad
+H_{l+1}=A_l+\operatorname{MLP}_{384\to1152\to384}(\operatorname{LN}(A_l)).
+$$
+
+The feed-forward activation is GELU and attention dropout is zero. These bridge
+blocks have no additional target-query self-attention layer. CLIP has already
+contextualized the source states. Output features and mask logits are:
+
+$$
+\widehat Y_j=\mu_j+\sigma\odot(W_oH_{4,j}+b_o),\qquad
+\ell_j=w_m^TH_{4,j}+b_m.
+$$
+
+The saved target mean $\mu$ is per slot and channel (77 × 2304); the standard
+deviation $\sigma$ is per channel (2304). Both come from teacher-valid training
+positions and were retained from the parent; standard deviation has a floor of
+0.1. The linear 384 → 2304 output head retains the old per-slot affine rank
+restriction of at most 384. Winning validation does not imply that this model
+incorporates every architectural repair or can represent arbitrary target states.
+
+This transfers the donor's text representation. The SD1.5 diffusion U-Net never
+runs in this path, so the bridge does not establish transfer of its generated
+style or composition. Spatial donor-feature transfer is a separate interface.
+
+## Exact active loss
+
+The winning objective is:
+
+$$
+\boxed{L=0.1L_{embed}+L_{functional}.}
+$$
+
+For native teacher conditioning $Y$, teacher-valid mask $M$, and bridge output
+$\widehat Y$, embedding loss is masked, channel-normalized squared error:
+
+$$
+L_{embed}=\frac{\sum_{b,j,c}M_{b,j}
+\left((\widehat Y_{b,j,c}-Y_{b,j,c})/\sigma_c\right)^2}
+{2304\max(1,\sum_{b,j}M_{b,j})}.
+$$
+
+Functional distillation asks whether the **same frozen SANA** makes the same
+denoising prediction under bridged and native conditioning. Student predictions
+carry gradients through SANA into the bridge. Teacher predictions are detached
+or read from the native-state cache; SANA's parameters remain frozen.
+
+At the same latent state and timestep, let $s_+,s_0$ denote SANA predictions
+under student positive/empty conditioning, and $t_+,t_0$ its teacher predictions.
+All MSEs average across the functional minibatch and latent coordinates:
+
+$$
+g_s=s_0+4.5(s_+-s_0),\qquad g_t=t_0+4.5(t_+-t_0),
+$$
+
+$$
+L_{functional}=\frac{
+\operatorname{MSE}(s_+,t_+)+\operatorname{MSE}(s_0,t_0)
++0.05\operatorname{MSE}(g_s,g_t)}
+{\max(10^{-5},\operatorname{mean}((t_+-t_0)^2))}.
+$$
+
+The denominator is detached **native prompt-effect power**. It measures error
+relative to how strongly the teacher responds to the prompt. The positive
+student uses its predicted hard mask; the teacher uses its cached native mask.
+Their mismatch matters when predicted slots are wrong.
+
+Both empty branches use native empty conditioning. When their predictions agree
+exactly, empty error vanishes and the guided error becomes $4.5^2$ times positive
+error. The numerator then simplifies to **2.0125 times positive MSE**. The code
+retains the full expression, including any cached-versus-live numerical
+discrepancy. This run principally matches the conditional recipient response;
+it does not learn a new unconditional branch.
+
+LPIPS, pixel error, and CLIP judge rendered outputs for selection/evaluation;
+they are not differentiable image losses in this run. Full-sampler
+backpropagation was tested in separate arms and did not win the combined
+validation ranking.
+
+## What training changes
+
+The run initializes the already-trained parent and updates the **887,040 output
+projection parameters** that receive nonzero gradients. The `heads` setting also
+enables 385 mask-head parameters, but mask BCE has weight zero and hard thresholds
+supply no functional mask gradient. Their values remain unchanged during training.
+
+| Setting | Winning continuation |
+|---|---|
+| Optimizer | Adam, betas (0.0, 0.999), epsilon 1e-8, no weight decay |
+| Learning rate | 1e-5; held through 60% of the run, then cosine toward a 5% floor |
+| Updates and selection | 1,500 completed; EMA update 1,250 selected |
+| EMA | Decay 0.98 |
+| Gradient clipping | Global bridge gradient norm at 1.0 |
+| Embedding batch | 32 rows: 31 sampled captions plus one empty row |
+| Functional batch | Four positive examples every update |
+| State mixture | Each update chooses cached native states with probability 0.5, otherwise fresh initial noise |
+| Native cache | 512 training captions × five states, with cached teacher predictions |
+| Training data | 8,192 paired captions; no image reconstruction cache in this run |
+| Frozen components | Source projection, queries, four blocks, router, cloud, particle projection, calibration buffers, CLIP, SANA, DC-AE |
+| Inactive loss terms | GAN/b_cap, critic features, mask BCE, particle VICReg/centering, padding, soft masks, full rollouts |
+
+After training, mask calibration **subtracts 2.0417721271514893** from the old
+mask bias, equivalent to raising the old-logit threshold by that amount. It
+minimizes mask errors on 58 validation prompts. Ordinary inference thresholds
+the resulting logits at zero and always preserves the first BOS slot.
+Calibration changes one scalar and requires no fitting during deployment.
+
+Checkpoint selection uses rendered validation LPIPS and is separate from the
+optimized loss. The release contains EMA weights, not the last live optimizer
+iterate. Final testing occurred after freezing selection.
+
+The observed improvement is **1.42% lower held-out LPIPS and 2.50% lower pixel
+MAE**, with confidence intervals including no improvement and composition errors
+remaining. Calibration adds two missed test positions. A nonlinear frozen
+recipient can respond better to an improved linear head without that head
+recovering all inaccessible target directions.
+
+## Expanded-source alternative
+
+The alternative preserves the parent and adds this branch to normalized output:
+
+$$
+z_j=W_oH_{4,j}+b_o+
+W_2\operatorname{SiLU}(W_1H_{4,j}+b_1+W_sX_j)+b_2.
+$$
+
+$W_1$ maps 384 → 2304, $W_s$ maps the full 768 source channels directly to 2304
+without bias, and $W_2$ maps 2304 → 2304. Source information can bypass the old
+768 → 384 projection, and the nonlinear branch removes its affine-output
+restriction. This does not promise arbitrary mappings or aligned word positions.
+
+This model has **15,150,149 parameters**, including **7,967,232 added branch
+parameters** trained with the same active loss and schedule. The last added
+projection starts at zero to preserve parent predictions; the source projection
+uses full fan-in initialization scale 1.0. Its selected EMA update is 1,000,
+followed by the same mask calibration. It improves mean test fidelity but loses
+the combined validation ranking to the linear continuation. Both branches run
+within one ordinary bridge forward.
+
+## Deployment and source of truth
+
+The deployment CLI uses CLIP FP16 with batches padded to 16 to match the cached
+encoder kernel shape, a FP32 bridge, and BF16 SANA. It supplies predicted positive
+conditioning and the stored native empty conditioning to the pinned 20-step DPM
+sampler at guidance 4.5 and 512 × 512. The bridge output is computed once and
+reused throughout sampling; the original DC-AE decodes the final latent. No Gemma
+encoder, teacher cache, retrieval, optimizer, or best-of-N is needed by inference.
+
+The equations follow [the bridge](https://github.com/255BITS/model-glue/blob/main/model_glue/text_bridge.py),
+[training and guidance metrics](https://github.com/255BITS/model-glue/blob/main/model_glue/text_train.py),
+[inference](https://github.com/255BITS/model-glue/blob/main/model_glue/sana_infer.py), and
+[mask calibration](https://github.com/255BITS/model-glue/blob/main/scripts/calibrate_text_mask.py). Exact model/checkpoint
+revisions are recorded in the release manifest and experiment JSON.

README.md

@@ -13,6 +13,7 @@ tags:
 base_model:
 - stable-diffusion-v1-5/stable-diffusion-v1-5
 - Efficient-Large-Model/Sana_600M_512px_diffusers
+base_model_relation: adapter
 inference: false
 ---
 
@@ -42,6 +43,68 @@ retains the linear head. All three checkpoints include the complete bridge, clou
 calibration buffers, architecture metadata, and pinned source/recipient identities.
 They load directly without the parent checkpoint or training caches.
 
+## Architecture and winning objective
+
+The default contains **7,182,917 parameters**. SD1.5 CLIP provides 77 × 768
+contextual states. LayerNorm and a linear 768 → 384 layer produce source features.
+A mean-pooled source passes through a 768 → 64 → 4 SiLU router; softmax attention
+over the saved 256 × 4 particle cloud gives a deterministic four-dimensional
+mixture. Its projection offsets 77 learned queries, which also receive the
+projected source states.
+
+Four width-384 blocks each apply six-head cross-attention to the source followed
+by a 384 → 1152 → 384 GELU feed-forward network, with residual connections and
+LayerNorm. A final linear 384 → 2304 head produces normalized conditioning,
+rescaled with saved target statistics. A separate linear 384 → 1 head predicts
+mask logits. The first BOS slot is always kept.
+
+The donor diffusion U-Net is outside this path. This checkpoint transfers CLIP
+text states; it does not establish transfer of SD's generated style or composition.
+
+The winning continuation updates only the **887,040 output-projection parameters**
+that receive nonzero gradients. The parent stem, router, cloud, and mask values
+remain fixed during training. Afterward, validation calibration subtracts
+2.0417721271514893 from the old mask-head bias. The release stores that fixed
+adjustment and thresholds its logits at zero.
+
+The active training objective is:
+
+$$
+L=0.1L_{embed}+
+\frac{\operatorname{MSE}(s_+,t_+)+\operatorname{MSE}(s_0,t_0)
++0.05\operatorname{MSE}(g_s,g_t)}
+{\max(10^{-5},\operatorname{mean}((t_+-t_0)^2))},
+$$
+
+where $L_{embed}$ is embedding squared error on teacher-valid tokens, divided by
+the saved per-channel variance. $s_+,s_0$ and $t_+,t_0$ are frozen-SANA predictions
+with student and native positive/empty conditioning at the same latent and
+timestep. The guided predictions are $g_s=s_0+4.5(s_+-s_0)$ and
+$g_t=t_0+4.5(t_+-t_0)$. The denominator is detached native prompt-effect power.
+Student calls backpropagate into the bridge; teacher calls and SANA parameters
+remain frozen. Both empty branches use stored native conditioning.
+
+The run uses Adam at initial LR 1e-5, betas (0.0, 0.999), delayed cosine decay,
+gradient clipping at 1.0, and EMA 0.98. Each update uses 32 embedding rows and
+four functional examples; it chooses cached native states with probability 0.5
+and fresh initial noise otherwise. Selection uses rendered validation LPIPS.
+
+**GAN, b_cap, critic-feature, particle, mask-BCE, soft-mask, and full-rollout losses
+are inactive in this winning continuation.** Its parent has ParticleGAN provenance.
+LPIPS and CLIP are evaluation signals here, not differentiable image losses.
+The bridge runs once per prompt; ordinary SANA sampling reuses its conditioning.
+
+The expanded-source alternative adds
+$W_2\operatorname{SiLU}(W_1h+W_sx+b_1)+b_2$ alongside the original output head.
+Its hidden branch is 2304-wide and sees the full 768-channel source directly.
+It has **15,150,149 parameters**, with **7,967,232 new branch parameters** trained
+under the same active loss. It removes the old source-projection blind space and
+affine-output restriction but did not win combined validation.
+
+The [complete formulation](https://huggingface.co/ntc-ai/model-glue-sd15-sana-text/blob/main/FORMULATION.md)
+includes equations, an architecture diagram, parameter counts, the mask
+calibration sign, training controls, and deployment details.
+
 ## Run
 
 Use the implementation revision that was tested for this release:

SHA256SUMS

@@ -1,5 +1,6 @@
+cd4bee587e4e1d8dc12049bccc5ce11e8aba5f1c3d3de74b160d67a1535562e8  FORMULATION.md
 f50d7a3340d8e33d65e6905d6471ee64c2217a0a20f54a67922f96e02ec6a42b  LICENSE
-8786e5eaf37799205bbc1f25cdd4cf4d965a9ad208e544e3a53d2656f7ae62a6  README.md
+749dc3da10cbb821da5a47641e047fe85192ffa696939865b9f6451e0cb01c5c  README.md
 1a2db74c6c2d940ee15a845a7e60404ffbe32573757672586a18c578db36a10a  alternatives/expanded-source.safetensors
 fc7f44369cf21791c4f9ad6b74978dd420be41a76330a1fe32b11f440f3bf925  alternatives/linear-uncalibrated.safetensors
 aea09893f34769c9b4eef9c30ab9711c935854ec3cf29571aee81cc2367109d7  config.json
@@ -10,7 +11,7 @@ ceca1ec0fcac308fc6a565452f3dc67ce90dcd21e323acefd6df21cc9c1b78b4  configs/source
 4dc2536374ff7b33c85e7bf21e55977737d4e3cd7d0a6db9189812972dbdf469  configs/validation-confirmation.json
 12de8ae1a69eff31879911c50068af1a2fe77bd5f56dbbf5cf1e6743464eca06  figures/bench-comparison.jpg
 d3f12c915d977e11bbffc83b6e1ba94fabe5feafe2a3a35b0ee6c84e34e77654  figures/first-four-test-prompts.jpg
-635da76e075a43329ac4700a2f3a0e5cb71434be04f0828de29fc0a5b69acbcb  manifest.json
+97e480871c159c26545f77c87bc81d6c9b83e55bd5817c6beba16d9701f6656a  manifest.json
 d3be2492e49f6eb1bc28a406d35108f1055e3dd19e412792442b4c5d9274bb4d  model.safetensors
 5416eac9605b2e0af6b5325be168dbd86bdb8a7d5f241f5481dee417e439c349  reports/mask-metrics.json
 fabd67b38e5a5e878f63bbb0f6de0278c4bcc459d0603957ce3603417f1fddf2  reports/retraining.json

manifest.json

@@ -142,13 +142,17 @@
     }
   },
   "files": {
+    "FORMULATION.md": {
+      "sha256": "cd4bee587e4e1d8dc12049bccc5ce11e8aba5f1c3d3de74b160d67a1535562e8",
+      "size_bytes": 10246
+    },
     "LICENSE": {
       "sha256": "f50d7a3340d8e33d65e6905d6471ee64c2217a0a20f54a67922f96e02ec6a42b",
       "size_bytes": 1064
     },
     "README.md": {
-      "sha256": "8786e5eaf37799205bbc1f25cdd4cf4d965a9ad208e544e3a53d2656f7ae62a6",
-      "size_bytes": 7968
+      "sha256": "749dc3da10cbb821da5a47641e047fe85192ffa696939865b9f6451e0cb01c5c",
+      "size_bytes": 11406
     },
     "alternatives/expanded-source.safetensors": {
       "sha256": "1a2db74c6c2d940ee15a845a7e60404ffbe32573757672586a18c578db36a10a",