LookWhere? model scaling and deployment

Vision transformer training on Cloud TPU pods (1 TB HBM, 10 petaFLOPs FP16): mesh topology, per-host utilization, and honest throughput numbers. NeurIPS 2025.

NeurIPS 2025

LookWhere? Efficient Visual Recognition by Learning Where to Look and What to See from Self-Supervision

Anthony Fuller, Yousef Yassin, Junfeng Wen, Tarek Ibrahim, Daniel G. Kyrollos, James R. Green, Evan Shelhamer

Advances in Neural Information Processing Systems (NeurIPS), 2025, vol. 38, pp. 96710–96748

BibTeX

@inproceedings{fuller2025lookwhere,
  title         = {{LookWhere}? Efficient Visual Recognition by Learning Where to Look and What to See from Self-Supervision},
  author        = {Anthony Fuller and Yousef Yassin and Junfeng Wen and Tarek Ibrahim and Daniel G. Kyrollos and James R. Green and Evan Shelhamer},
  booktitle     = {Advances in Neural Information Processing Systems (NeurIPS)},
  year          = {2025},
  volume        = {38},
  pages         = {96710--96748},
  doi           = {10.52202/085713-2910},
  eprint        = {2505.18051},
  archivePrefix = {arXiv},
  url           = {https://neurips.cc/virtual/2025/loc/san-diego/poster/117306}
}

LookWhere? (Fuller, Yassin, Wen, Ibrahim, Kyrollos, Green, Shelhamer; NeurIPS 2025) learns where to look and what to see from self-supervision, so a vision transformer spends its compute on the parts of an image that matter. My part was the systems side: making the models train and evaluate efficiently on Cloud TPU pods from the Google TPU Research Cloud.

On a pod with four chips per host in megacore mode, mesh topology and per-host utilization decided throughput far more than model design. Most of my time went into the sharding layout and into measuring what the hardware was actually doing, which is harder than it sounds. I once discarded a run after finding an error in how per-device FLOP rates were aggregated: the number looked excellent and was wrong. Since then I treat a throughput figure I have not checked as a claim, not a result.

The pod had on the order of 1 TB of high-bandwidth memory and 10 petaFLOPs of FP16 compute. Keeping it busy was the job.