vit-ae-2x-f8c32

<img src="https://cdn-uploads.huggingface.co/production/uploads/608c79739c5a8f8ddd85c409/_ENfOJ_3kNxanMEpIdnx8.gif" width="60%">

commaai/vit-ae-2x-f8c32

Authors: Armand du Parc Locmaria and Yassine Yousfi. comma.ai

This ViT autoencoder is trained on paired narrow and wide view camera images captured from comma.ai devices. Each image is calibrated and resized to (3x128x256); vertically concatenated to a (3x256x256) input, which is compressed to a (32x16x32) latent tensor using (16x8) patches.

The encoder and decoder contain 6 and 12 ViT blocks, respectively. Both use width 768, 6 attention heads, a 4×GELU MLP, QK normalization, fixed 2-D sinusoidal positional embeddings. The 512 width-768 patch tokens are projected to 32 latent channels. The encoder and decoder have 42.84M and 85.35M parameters, respectively.

Training combines mean-squared error (MSE), Learned Perceptual Image Patch Similarity (LPIPS), and adversarial loss.

Training proceeds in two stages:

1. MAE pretraining (132M paired-view samples): jointly train the encoder and decoder with 75% random patch masking using MSE and LPIPS, without adversarial loss [1].

2. Adversarial decoder fine-tuning (26.7M paired-view samples): remove masking, freeze the encoder, and train the decoder adversarially alongside the discriminator.

The discriminator follows the paired Reconstruction GAN (rGAN) design [2], where the target and reconstruction are concatenated channel-wise.

The 0.70M-parameter discriminator (PatchGAN) uses no gradient penalty or R1 regularization. It is intentionally kept small, trained with a small learning rate and a large AdamW weight decay.

Reconstruction quality

We report five reconstruction metrics:

• LPIPS measures per-image perceptual reconstruction distance; lower is better.
• rFID compares the Inception-feature distributions of original and reconstructed images; lower is better.
• Sharpness is the mean grayscale Sobel-gradient magnitude; higher values indicate stronger local intensity gradients.
• YOLO confidence averages post-NMS objectness times maximum class probability over retained detections, an indirect measure of recognizable object content.
• Mean velocity error is the mean error between the predicted velocities for original and reconstructed images using a PoseNet, a proxy for temporal fidelity.

Modelability

Reconstruction quality alone does not determine whether a latent space is useful for generative modeling; its distribution must also be learnable by a prior [3].

For each representation, we train the same DiT-L/2 architecture [4] (24 blocks, width 1024, 16 heads, 2×2 latent patches) and compute FID on generated images (gFID).

Training uses a rectified-flow velocity objective with log-normal noise-level sampling; generation uses 50 Euler steps.

<img src="https://cdn-uploads.huggingface.co/production/uploads/608c79739c5a8f8ddd85c409/SbsRsqQKvwuejFCVihRul.png" width="60%">

Usage

See example.py for how to use this with the commaai/comma1M dataset

References

[1] Chen, H., Han, Y., Chen, F., Li, X., Wang, Y., Wang, J., Wang, Z., Liu, Z., Zou, D. and Raj, B., 2025, January. Masked autoencoders are effective tokenizers for diffusion models. In Proceedings of the Forty-second International Conference on Machine Learning.

[2] HaCohen, Y., Chiprut, N., Brazowski, B., Shalem, D., Moshe, D., Richardson, E., Levin, E., Shiran, G., Zabari, N., Gordon, O., Panet, P., Weissbuch, S., Kulikov, V., Bitterman, Y., Melumian, Z. and Bibi, O., 2024, December. LTX-Video: Realtime Video Latent Diffusion. arXiv preprint arXiv:2501.00103.

[3] Dieleman, S., 2025, April. Generative modelling in latent space.

[4] Peebles, W. and Xie, S., 2023. Scalable Diffusion Models with Transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision.