CanViT: Toward Active-Vision Foundation Models (NeurIPS 2026)

Yohaï-Eliel Berreby^1,2 · Sabrina Du^1,2 · Audrey Durand^2,3 · B. Suresh Krishna^1
1McGill University · ²Mila – Quebec AI Institute · ³Université Laval

This repository holds the reference PyTorch implementation of CanViT, the Canvas Vision Transformer, with pretraining, task specialization and evaluation; the package is canvit-pytorch on PyPI. The native MLX and JAX/Flax NNX packages live under canvit-mlx/ and canvit-nnx/ and share canvit-core/.

News

• 2026-09-26: canvit-pytorch 0.2, a refactoring release; updated checkpoints pushed to the Hub. Code written for 0.1: see Troubleshooting.
• 2026-09-24: 🎉 Accepted at NeurIPS 2026!
• 2026-05-16: Preprint v2 (arXiv:2603.22570v2), adding the 84.5% ImageNet-1k fine-tuning result and the effect of canvas resolution.
• 2026-04-06: First finetuned IN1k checkpoint: canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-04-06, with new CanViTForImageClassification API.
  • 🎉 CanViT sets a new SOTA on active-vision IN1k classification, with 84.5% top-1 accuracy, up from AdaptiveNN's previous best of 82.2%.
• 2026-03-23: Preprint v1 (arXiv:2603.22570v1).
  • 🎉 CanViT sets a new SOTA on active ADE20K segmentation, with 45.9% ADE20K mIoU, obtained using linear probing from frozen weights.
• 2026-02-18: Initial code and first pretrained checkpoint release.

Abstract

Active computer vision promises efficient, biologically plausible perception through sequential, localized glimpses, but lacks scalable general-purpose architectures and pretraining pipelines, leaving Active-Vision Foundation Models (AVFMs) underexplored. We introduce CanViT, the first task- and policy-agnostic AVFM. CanViT uses scene-relative RoPE to bind a retinotopic Vision Transformer backbone and a spatiotopic scene-wide latent workspace, the canvas. Efficient interaction with this high-capacity working memory is supported by Canvas Attention, a novel asymmetric cross-attention mechanism. We decouple thinking (backbone-level) and memory (canvas-level), eliminating canvas-side self-attention and fully-connected layers to achieve fast sequential inference and scalability to high output resolutions. We propose a label-free active vision pretraining scheme, policy-agnostic passive-to-active dense latent distillation: reconstructing scene-wide DINOv3 embeddings from sequences of low-resolution glimpses with randomized locations, zoom levels, and lengths. We pretrain CanViT-B from a random initialization on 13.2 million ImageNet-21k scenes—an order of magnitude more than previous active models—and 1 billion random glimpses, in 166 hours on a single H100. On ADE20K segmentation, a frozen CanViT-B achieves 38.5% mIoU in a single low-resolution glimpse, outperforming the best active model's 27.6% with 20x fewer inference FLOPs as well as its FLOP- or input-matched DINOv3 teacher. Given additional glimpses, CanViT-B reaches 45.9% ADE20K mIoU. On ImageNet-1k classification, CanViT-B also sets a new active-vision state of the art, with 84.5% top-1 accuracy after fine-tuning. CanViT generalizes to longer rollouts, larger scenes, and new policies. Our work narrows the wide gap between passive and active computer vision, demonstrating the potential of task- and policy-agnostic AVFM pretraining.

Checkpoints

The table lists PyTorch checkpoints on Hugging Face. Native MLX and JAX/Flax NNX checkpoints are linked in Backend packages.

Checkpoint	Description
canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02	CanViT-B pretrained on ImageNet-21k by policy-agnostic passive-to-active dense latent distillation from DINOv3 ViT-B: the paper's model
canvitb16-add-vpe-pretrain-g128px-s512px-in1k-dv3b16-2026-06-22	CanViT-B pretrained the same way on ImageNet-1k only
canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-04-06	The ImageNet-21k model fine-tuned for ImageNet-1k classification, 84.5% top-1 (trained on TPU v6e via torch_xla)
canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-07-24	The ImageNet-21k model fine-tuned again with the same recipe, by a JAX/Flax NNX trainer
canvitb16-add-vpe-finetune-g128px-s512px-in1k-from-in1k-2026-07-24	The ImageNet-1k model fine-tuned with the same recipe, by the same trainer

Collections:

• CanViT ADE20K segmentation probes and DINOv3 ADE20K segmentation probes
• DINOv3 ImageNet-1k probes, which the frozen classification path uses
• CanViT pretraining ablations, the paper's ablation backbones
• JAX / Flax NNX and MLX copies, for the ports under Repository layout

Quickstart

We recommend uv for dependency management.

uv add canvit-pytorch

The snippets use path/to/image.jpg; replace it with the path to an RGB image.

from canvit_pytorch import CanViTForPretraining, Viewpoint, sample_at_viewpoint
from canvit_pytorch.preprocess import preprocess
from PIL import Image
import torch

# CanViT is integrated with the HuggingFace Hub.
model = CanViTForPretraining.from_pretrained(
    "canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02"
).eval()

image = Image.open("path/to/image.jpg").convert("RGB")
image = preprocess(512)(image)
image = image.unsqueeze(0)  # [1, 3, 512, 512]

# CanViT is a recurrent model.
state = model.init_state(batch_size=1, canvas_grid_size=32)

with torch.inference_mode():
    vp = Viewpoint.full_scene(batch_size=1, device=image.device)
    glimpse = sample_at_viewpoint(spatial=image, viewpoint=vp, glimpse_size_px=128)
    out = model(glimpse=glimpse, state=state, viewpoint=vp)

# The output carries the updated recurrent state and the glimpse features.
out.state.recurrent_cls  # [1, 1, 768] — global CLS token
out.glimpse_patches      # [1, 64, 768] — glimpse patch features

# Now let's do a second glimpse: zoom into the top-left quadrant
# You can do this repeatedly: CanViT is recurrent with a large but constant-size canvas.
with torch.inference_mode():
    vp2 = Viewpoint(
        centers=torch.tensor([[-.5, -.5]], device=image.device),
        scales=torch.tensor([.5], device=image.device),
    )
    glimpse2 = sample_at_viewpoint(spatial=image, viewpoint=vp2, glimpse_size_px=128)
    out2 = model(glimpse=glimpse2, state=out.state, viewpoint=vp2)
    

ImageNet-1k Classification

CanViTForImageClassification provides a unified interface for classification. Two construction paths, same forward pass:

From a finetuned checkpoint (CanViT + head trained on IN1k):

from canvit_pytorch import CanViTForImageClassification, Viewpoint, sample_at_viewpoint
from canvit_pytorch.preprocess import preprocess
from PIL import Image
import torch

clf = CanViTForImageClassification.from_pretrained(
    "canvit/canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-04-06"
).eval()

From the frozen pretrained CanViT checkpoint + a DINOv3 linear probe:

clf = CanViTForImageClassification.from_pretrained_with_probe(
    pretrained_repo="canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02",
    probe_repo="canvit/dinov3-vitb16-lvd1689m-in1k-512x512-linear-clf-probe",
).eval()

Both have the same forward pass:

image = preprocess(512)(Image.open("path/to/image.jpg").convert("RGB")).unsqueeze(0)
state = clf.init_state(batch_size=1, canvas_grid_size=32)

with torch.inference_mode():
    vp = Viewpoint.full_scene(batch_size=1, device=image.device)
    glimpse = sample_at_viewpoint(spatial=image, viewpoint=vp, glimpse_size_px=128)
    logits, state = clf(glimpse=glimpse, state=state, viewpoint=vp)

print(logits.argmax(dim=-1))  # ImageNet-1k class index

ADE20K Semantic Segmentation

CanViTForSemanticSegmentation bundles a CanViT and a SegmentationProbe head into one model. forward returns per-pixel logits at canvas-grid resolution; predict adds bilinear upsampling.

import torch
from PIL import Image

from canvit_pytorch import CanViTForSemanticSegmentation, Viewpoint, sample_at_viewpoint
from canvit_pytorch.preprocess import preprocess

seg = CanViTForSemanticSegmentation.from_pretrained_with_probe(
    pretrained_repo="canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02",
    probe_repo="canvit/probe-ade20k-40k-s512-c64-in21k",
).eval()

image = preprocess(512)(Image.open("path/to/image.jpg").convert("RGB")).unsqueeze(0)
state = seg.init_state(batch_size=1, canvas_grid_size=64)
with torch.inference_mode():
    viewpoint = Viewpoint.full_scene(batch_size=1, device=image.device)
    glimpse = sample_at_viewpoint(spatial=image, viewpoint=viewpoint, glimpse_size_px=128)
    upsampled, state = seg.predict(
        glimpse=glimpse,
        state=state,
        viewpoint=viewpoint,
        target_size=(1024, 1024),
    )  # [B, n_cls, 1024, 1024]

The standalone SegmentationProbe head is also exported from canvit_pytorch for use on any spatial feature map. The published probes are listed under Checkpoints.

Backend packages

CanViT's PyTorch, MLX and JAX/Flax NNX implementations share model names and architecture definitions in canvit-core/. Constructing an NNX model from scratch also takes rngs=nnx.Rngs(0).

uv add canvit-mlx  # MLX on Apple Silicon
uv add canvit-nnx  # JAX / Flax NNX

Native arrays use NHWC: a model glimpse has shape [B, H, W, 3], and sample_at_viewpoint accepts [B, H, W, C]. PyTorch tensors use NCHW: images and feature maps are [B, C, H, W], while model glimpses are [B, 3, h, w]. Viewpoint centers are (row, col) scene coordinates in [-1, 1], and scales are crop half-sides across all backends.

The MLX README and NNX README contain installed-package examples that load full Hugging Face IDs. Native models and segmentation probes are separate checkpoints; the package READMEs link the available model and probe choices.

Demos

The demos are repository files. From a checkout, run them in canvit-pytorch/:

cd canvit-pytorch

# Classification with sequential glimpses
uv run --extra demo python demos/classify.py                     # finetuned checkpoint
uv run --extra demo python demos/classify.py --classifier frozen  # frozen CanViT + fused probe

# Canvas PCA visualization with two viewing strategies
uv run --extra demo python demos/basic.py

Supported platforms

• CPU
• CUDA (tested on RTX 4090, H100 SXM 80GB)
• Apple Silicon via MPS (inference, tested with torch 2.14)
• TPU via torch_xla 2.9.0 (tested on TPU v6e)

We aim to maintain compatibility with torch.export and ONNX Runtime. Please file an issue if you encounter problems.

Repository layout

canvit-pytorch/ holds the package:

Module	Contents	Docs
canvit_pytorch	The model: architecture, checkpoint loading, viewpoints and policies, task heads	this README
canvit_pytorch.flops	Analytic forward FLOPs of CanViT and DINOv3	docs/flops.md
canvit_pytorch.pretrain	Pretraining: passive-to-active dense distillation from DINOv3	docs/pretrain.md
canvit_pytorch.specialize	Downstream training: ADE20K segmentation probes, ImageNet-1k fine-tuning	docs/specialize.md
canvit_pytorch.evaluate	Evaluation and benchmarking: ADE20K mIoU, ImageNet-1k top-k, DINOv3 reconstruction	docs/evaluate.md
canvit_pytorch.viz	Recorded rollouts and smooth viewpoint paths for the project page and slides	docs/viz.md
canvit-core	Shared configuration, architecture specifications, readout fusion and checkpoint schema	README
canvit-mlx	MLX implementation for Apple Silicon	README
canvit-nnx	JAX/Flax NNX implementation	README

site/ is the project page.

paper/ holds the NeurIPS 2026 paper: its LaTeX sources, and the pipeline that turns evaluation results into its tables and figures.

Related repositories:

• dinov3-in1k-probes — ImageNet-1k linear probes for the DINOv3 ViTs, used by CanViT's classification path

Troubleshooting

If you encounter errors loading pretrained checkpoints, ensure you are using the latest version of the package:

uv lock --upgrade-package canvit-pytorch && uv sync

Code written for canvit-pytorch 0.1 needs canvit-pytorch<0.2 and the checkpoints' earlier files: from_pretrained(..., revision="canvit-pytorch-0.1").

Citation

If you use this work, please cite our paper:

@article{berreby2026canvit,
  title={CanViT: Toward Active-Vision Foundation Models},
  author={Berreby, Yoha{\"i}-Eliel and Du, Sabrina and Durand, Audrey and Krishna, B. Suresh},
  year={2026},
  eprint={2603.22570},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2603.22570}
}

Contact

Open an issue in this repository or email me@yberreby.com.

Star History

License

MIT. See LICENSE for details.