Add Pixio

[lorinczszabolcs]

This PR adds Pixio to lightly.

Pixio is a masked-autoencoder method: it builds on MAE with three changes:

• a deeper decoder

• coarser block (grid) masking

• multiple class tokens whose mean is the global representation. This is a reimplementation from the paper; the reference repository is under the FAIR Noncommercial license, so no code or weights are taken.

• Paper: In Pursuit of Pixel Supervision for Visual Pre-training, CVPR 2026

• Reference implementation: https://github.com/facebookresearch/pixio

Closes #1894.

What's included

• lightly/models/utils.py — random_grid_token_mask, block/grid masking returning lightly's (idx_keep, idx_mask) convention, with input validation (positive patch count, perfect-square grid, block divisibility). Also generalizes the 2D sincos positional-embedding helpers to N prefix tokens (fixing a latent bug where >1 prefix token was mis-sized).

• lightly/models/modules/masked_autoencoder_timm.py — PixioDecoderTIMM (a thin MAEDecoderTIMM subclass, default decoder depth 32) plus a num_prefix_tokens parameter on MAEDecoderTIMM (default 1, backward-compatible).

• benchmarks/imagenet/vitb16/pixio.py — ViT-B/16 benchmark (256px, 4×4 grid mask, 8 class tokens via reg_tokens=7, 512×32 decoder, normalized-pixel loss), registered in main.py.

• PyTorch / PyTorch-Lightning / Lightning-Distributed examples (+ generated notebooks) and an example docs page, added to the models list. PDF attached. pixio.pdf

• Tests in tests/models/test_ModelUtils.py and tests/models/modules/test_masked_autoencoder_timm.py.

Testing

• New tests cover the mask util (keep/mask partition, whole-cell granularity, keep-count, mask_ratio 0.0/1.0 extremes, grid_size=4, error paths, device) and the decoders (CPU + CUDA forward, output shapes, finite values).
• make format, mypy, ruff, and the Pixio tests pass locally. make all-checks is green except three pre-existing, environment-only failures unrelated to this change (a distributed test_loss_dcl and two test_version_checking network-timeout tests).
• Ran a short training of the model on FashionMNIST with a shallower network (decoder depth=2) due to hardware constraints: the loss decreased monotonically. A full ImageNet benchmark isn't possible on my side, the PR just contributes the method, not a SOTA reproduction (the paper's numbers rely on web-scale data + distillation).

Notes

• 256px / patch-16 so the 16×16 patch grid divides evenly into 4×4 mask blocks; docs note the 2×2-grid / 4-class-token configuration reported as dense-optimal in the paper's ablations.
• Relationship to Add random block wise mask from Pixio #1919: that PR independently adds a block-wise mask for this issue. Its current implementation lifts noise[:, 0] = -1 from random_token_mask, where that line protects the CLS token at index 0, but in block space index 0 is a spatial block, so it deterministically keeps the top-left block in every sample (when num_prefix_tokens > 0 and ≥1 block is kept), biasing the masking. I confirmed this by running that function: the top-left block is kept in 100% of samples, vs the uniform 25% ours produces on the same config; That PRs own tests don't catch it because they only assert the prefix token at index 0 is kept. The current random_grid_token_mask protects prefix tokens structurally (unconditional prepend), so image blocks are masked uniformly at random.
• Unrelated fix: seeded the pre-existing test_normalize_mean_var and loosened its variance atol to 1e-4. It was latently flaky, unseeded torch.rand plus a tolerance too tight for normalize_mean_var's eps regularization (output variance is 1 - eps/var, not exactly 1), and this PR's new tests shifted global RNG ordering, which surfaced it.

[gabrielfruet]

Thank you so much for the nice contribution! I'm going to review it ASAP.

[gabrielfruet]

/review

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86.60%. Comparing base (fc3df0d) to head (b07df8a).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1965      +/-   ##
==========================================
+ Coverage   86.54%   86.60%   +0.06%     
==========================================
  Files         174      174              
  Lines        7467     7501      +34     
==========================================
+ Hits         6462     6496      +34     
  Misses       1005     1005              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[lorinczszabolcs]

Regenerated the notebooks in the meantime, as that test has failed.

[gabrielfruet]

LGTM! Very nice work, just small comments and we are good to go!

[gabrielfruet]

Thank you for the really nice first contribution!! Welcome to the open-source world :)

I'll try to accomodate some time to run benchmarks for Pixio

[gabrielfruet]

I missed that we should also update the README.md. If you wish, you can open a follow-up PR for that. I would just say to wait for #1975 to get merged before opening the PR or branching from it, since it changes the order of the papers.