savi-pretrained-30slots

Stochastic Slot Attention from Video (SAVi) encoder with 30 slots, pretrained on LMDrive BEV rasterizations. Stage 1a of the object-centric LMDrive pipeline.

Overview

This is one of four model artefacts released alongside a thesis on replacing the multimodal sensor-fusion encoder in LMDrive with an object-centric encoder based on Stochastic Slot Attention from Video (SAVi), evaluated on the CARLA LangAuto benchmark. The four artefacts are:

• ni-eminen/savi-pretrained-30slots -- Stage 1a encoder
• ni-eminen/savi-finetuned-30slots -- Stage 1b encoder
• ni-eminen/object-centric-lmdrive-4qt -- Q-Former + LLM (4 query tokens)
• ni-eminen/object-centric-lmdrive-8qt -- Q-Former + LLM (8 query tokens)

Architecture

• architecture: StoSAVi
• num slots: 30
• slot size: 256
• input resolution: [192, 192]
• num attention iterations: 2
• clip length: 6
• kld method: var-0.01

Inputs

192x192 RGB BEV images normalized to [0, 1] via torchvision.transforms.ToTensor(), arranged as temporal clips of 6 frames.

Training

Trained from scratch on LMDrive BEV rasterizations using the SlotFormer training framework. Optimizer Adam with LR 1e-4 and cosine annealing (2.5% warmup), gradient clipping 0.05, batch size 16 per GPU across 4 A100-40GB GPUs. Training ran for 11 epochs (approximately 108 hours).

Evaluation

Evaluated on 72 held-out BEV frames (12 videos x 6 frames) from test_town of the LMDrive dataset.

Metric	Value
Reconstruction MSE (per-pixel, [0, 1] scale)	0.00116 ± 0.00027
Reconstruction SSIM	0.9079 ± 0.014
FG-ARI (%)	39.94 ± 21.08
mIoU (%)	6.30 ± 2.80

Usage

import torch
# Model class defined in slotformer/bev_slots/models
# (see https://github.com/ni-eminen/SAVi-for-LMDrive-BEV)
from slotformer.bev_slots.models import build_model
from slotformer.bev_slots.configs.savi_bev_params import SlotFormerParams

params = SlotFormerParams()
model = build_model(params)
ckpt = torch.load("model.pth", map_location="cpu")
model.load_state_dict(ckpt["state_dict"], strict=False)
model.eval()

# Input: [B, T=6, 3, 192, 192] float tensor in [0, 1]
out = model({"img": video})
# out['post_slots']:            [B, T, 30, 256]  slot representations
# out['post_recon_combined']:   [B, T, 3, 192, 192] reconstructed image
# out['post_masks']:            [B, T, 30, 1, 192, 192] slot masks

Citation

@mastersthesis{nieminen2026objcentric,
  title  = {Object-centric world models for language-instructed autonomous driving},
  author = {Nieminen, Matias},
  year   = {2026},
  school = {Aalto University}
}

Related repositories

• Encoder training code: https://github.com/ni-eminen/SAVi-for-LMDrive-BEV
• Instruction-fine-tuning + evaluation code: https://github.com/ni-eminen/lmdrive-fix