phumane/kv-vae-qwen3-4b

KV-cache VAE for Qwen3-4B

eval_step03915.pt is a transformer VAE that compresses the KV cache of a reasoning prefix of Qwen/Qwen3-4B (thinking mode) into a single latent z ∈ R^1024.

What the file holds: weights (model_state) and config (cfg, a TransformerVAEConfig) only, with no optimizer state.

How it was trained:

• 6 encoder and 6 decoder layers, 305M parameters, input window of 4,096 positions.
• 5 epochs on 100,342 greedy Qwen3-4B DeepMath reasoning traces, 3,915 steps at an effective batch of 128.
• The KV cache was computed in bf16 and rounded through fp16.

How to load it: use analysis/scripts/build_latent_store.load_vae from the vae-training branch of the training repo. That repo also provides the models/ package that the pickled config needs. Encode with encode_trajectory, or with analysis/fields4b/encode_field.py.

Final eval:

Metric	Result
Reconstruction relative error, d = 64 / 256 / 1024 / 2048	0.567 / 0.557 / 0.544 / 0.609
Per-token KL at 25 / 50 / 75% of the trace	0.87 / 2.56 / 4.34

Note: reconstruction is well short of the 1.5B model's VAE.