Semantic alignment and flux.2 comparison
Hello,
I was making some comparisons between this model, flux.2 and my own and was surprised to find that KVAE 2.0's PCA exhibits very strong semantic alignment despite not having a semantic loss in the paper.
Upon closer inspection the PCA is very close to flux.2 and the cosine similarity is extremely high, around 97%.
Some details and examples here , see the PCA tab
https://huggingface.co/spaces/data-archetype/dinac_ae_d2-flux2-kvae-comparison
Was there some kind of alignment loss that was used during part of the training to keep the latent spaces in line ?
Hello,
Thanks for the interesting analysis.
We plan to share more details in future work.
In the meantime, you may want to run the same analysis on model#2 as well.
It achieves comparable reconstruction and generative evaluation performance - see https://huggingface.co/kandinskylab/KVAE-2D-2.0#versions
about v2.1 for details.
The PCA is still very close to flux.2 on 2.1.
Cosine sim vs flux.2 is low initially but still 87% after a very brief fitting of a linear layer in between the two latent spaces so I am guessing that v2.1 is a finetune of v2.
Thanks for taking such a close look at this. Really nice investigation.
1. v2.1 is not a finetune of 2.0. The encoder was trained from scratch, and we plan to publish details in the future.
2. We re-ran your protocol as we understood it, and added REPA-E checkpoints as controls.
We kept your Space's geometry: each latent is folded 2x2 into H/16 x W/16 tokens, so FLUX.2 ends up with 128 channels. FLUX.2 is BN-whitened, the other side stays raw.
For data we took the first 1000 images of COCO val2017, center-cropped to 512, which gives 1.02M tokens. We then fit a least-squares map with no bias and measured mean cosine per token on tokens held out from the fit. The map only goes one way, from the FLUX.2 latent onto the other encoder's.
One note on the controls: REPA-E has no public f8 checkpoint with 32 channels, so we took the 16-channel ones. They come out at 64 channels instead of 128, which the map handles fine.
| Pair | Mean cosine after the linear map |
|---|---|
| FLUX.2 <-> REPA-E/e2e-sd3.5-vae | 0.917 |
| FLUX.2 <-> REPA-E/e2e-qwenimage-vae | 0.906 |
| FLUX.2 <-> KVAE-2D v2.1 | 0.842 |
| FLUX.2 <-> FLUX.1 | 0.779 |
Things get even more interesting with CKA (Centered Kernel Alignment), computed on the same tokens, and again on per-image averages, since CKA estimates no parameters from the data.
| Pair | per token CKA | per image CKA |
|---|---|---|
| FLUX.2 <-> REPA-E/e2e-sd3.5-vae | 0.832 | 0.885 |
| FLUX.2 <-> KVAE-2D v2.1 | 0.828 | 0.800 |
| FLUX.2 <-> REPA-E/e2e-qwenimage-vae | 0.789 | 0.848 |
| FLUX.2 <-> FLUX.1 | 0.263 | 0.407 |
What we see is that, on these metrics, the REPA-E checkpoints also end up with latents close to FLUX.2.
The REPA-E training code is public too, and here are the checkpoints: