Update weights for 2x1 to latest step

README.md

@@ -84,10 +84,10 @@ A diffusion model is trained against Gaussian noise, so the latent should look l
 KL term and a scaling factor someone measures after the fact, I put the statistics directly in the loss on the
 encoder output.
 
-- **Moment loss** (both): `mean^2 + (std - 1)^2`, weight 0.1. On `2x1` the mean and std are pooled over the
-  batch. On `2x2` they are per image and per channel, so every single image comes out mean 0 / std 1 in each
-  channel, not just the dataset on average.
-- **Gaussian order-statistic loss** (`2x2`): sort the latent values of each image and channel and MSE them
+- **Moment loss**: `mean^2 + (std - 1)^2`, weight 0.1, taken per image and per channel, so every single image
+  comes out mean 0 / std 1 in each channel, not just the dataset on average. The 2x1 spent its first 53k steps
+  with this pooled over the batch and switched to per image for the last 9k.
+- **Gaussian order-statistic loss**: sort the latent values of each image and channel and MSE them
   against the expected order statistics of a standard normal of the same size (Blom's approximation). This
   pins the whole distribution shape, not just the first two moments, so a flat gray image still encodes to
   something that looks like noise instead of a constant.
@@ -109,15 +109,15 @@ PSNR on these three 512 px images (dB):
 | image | Primal 2x1 | Flux.1 VAE | Primal 2x2 | Flux.2 VAE |
 |---|---|---|---|---|
 | img01 (people, studio) | 32.1 | 35.7 | 37.5 | 37.2 |
-| img02 (picnic, foliage) | 22.2 | 24.2 | 27.1 | 26.5 |
-| img03 (UI screenshot) | 21.0 | 24.7 | 27.4 | 25.9 |
+| img02 (picnic, foliage) | 22.0 | 24.2 | 27.1 | 26.5 |
+| img03 (UI screenshot) | 21.3 | 24.7 | 27.4 | 25.9 |
 
 The 2x2 matches or beats the Flux.2 VAE here at the same information budget (both are 32 values per 8×8 pixel
 block). Three images and PSNR alone are not a benchmark, and a KL VAE pays a reconstruction tax for its posterior
 noise that a deterministic one does not. The 2x1 trails the Flux.1 VAE on textured content. One channel at 2x
-is a 12:1 reduction with no slack, and it has about a third of the training steps of the 2x2. Its per-image
-latent std also drifts (0.9 to 1.4 here) because its moment loss was batch pooled. The per-image normalization
-on the 2x2 holds std within a few percent of 1.
+is a 12:1 reduction with no slack, and it has well under half the training steps of the 2x2 (62k vs 147k). Its
+per-image latent std still drifts (1.07 to 1.21 here) because it only got the per-image normalization for its
+last 9k steps. The 2x2, trained with it throughout, holds std within a few percent of 1.
 
 ## Training
 
@@ -135,13 +135,13 @@ The tricks that should make the latent space easier to diffuse:
 - **Noise-perturbed decode.** On 20% of steps the decoder gets `z + sigma * noise`, `sigma` drawn per image
   from `U(0, 0.5)`, and still has to reconstruct the clean image. The decoder learns to tolerate diffusion
   model error and the latent gets smooth: nearby codes decode to nearby images.
-- **Solid colors** (`2x2`): 5% of images are swapped for a random flat color. Flat regions are the hard case
-  for a per-image normalized latent.
-- **Critic** (`2x1`): a frozen SigLIP2 NaFlex ViT gives patch features from layers 4, 8 and 12 of the real and
-  reconstructed image at native resolution. Small spectral-normalized conv heads on those features train with
-  a hinge loss plus R1, and the decoder gets the generator term at weight 0.1, ramped in over the first 1000
-  steps. It sharpens the texture that MAE and LPIPS leave soft. The 2x2 was trained without a critic, it
-  seemed to converge fine without one.
+- **Solid colors**: 5% of images are swapped for a random flat color. Flat regions are the hard case for a
+  per-image normalized latent.
+- **Critic** (`2x1` only, last 9k steps): a frozen SigLIP2 NaFlex ViT gives patch features from layers 4, 8
+  and 12 of the real and reconstructed image at native resolution. Small spectral-normalized conv heads on
+  those features train with a hinge loss plus R1, and the decoder gets the generator term at weight 0.1,
+  ramped in over the first 1000 steps. It sharpens the texture that MAE and LPIPS leave soft. The 2x2 was
+  trained without a critic, it seemed to converge fine without one.
 
 Every step is one encode and one decode. On a flip/rotate step the transformed pair replaces the clean pair
 instead of adding a second decoder pass.

primal_vae_2x1.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9cf37dd0ff80edde481ee2cae858107a71855e6ed5fb495aa095a09b33b9189a
-size 321415508
+oid sha256:b035c96861e3d36ed0a5308e89535ba19c26846d98c358a3907949ff8d1776f3
+size 321415516