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

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

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 the Canvas Vision Transformer (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 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, listed by their Hub tag, for the ports under Repository layout

Quickstart

We recommend uv for dependency management. Add the package to your project (uv init creates one):

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

The demos run on the accelerator torch finds (CUDA, or MPS on Apple silicon), else on the CPU; --device chooses.

Supported platforms

• CPU
• CUDA (tested on RTX 4090, H100 SXM 80GB)
• Apple Silicon via MPS (inference, tested with torch 2.14)
• TPU: canvit-nnx for training and inference. canvit-pytorch runs inference under TorchTPU 0.2.0 (tested on TPU v5p); training under it is not supported yet.

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.rl	Learned viewing policies: run the published ones, train new ones on the frozen CanViT	docs/rl.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, ImageNet-1k fine-tuning with fused TPU kernels, and COCO object detection	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.

baselines/ holds the paper's baselines, apart from the packages: AdaGlimpse's architecture in JAX/Flax NNX, and the measurements of GFNet and AdaptiveNN on ImageNet-1k.

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. CanViT is pretrained by distillation from DINOv3, which is released under the DINOv3 License.