Portrait of Saurav Jha

Saurav Jha

saurav.jha@mila.quebec
Postdoctorate - Polytechnique Montréal
Supervisor
Sarath Chandar
Research Topics
Continual Learning
Generative Models
Website
Google Scholar
GitHub

Publications

recAMPLIFY: Large-Scale Pretraining of Recursive Protein Language Models
Saurav Jha
Qianli Liao
Robert Vernon
Sarath Chandar
Quentin Fournier
Protein language models (pLMs) are powerful tools for protein annotation and mutation-effect prediction, but progress has largely depended o… (see more)n scaling model size and data, making them expensive and prone to memorization. Test-time scaling offers a complementary path, but methods developed for autoregressive LLMs do not transfer directly: masked encoders lack a generation loop, and amino acids do not verbalize in a reasoning language. Instead, we adapt recursive pretraining to masked pLMs, where a weight-tied Transformer core refines residue-aligned hidden states over an adjustable number of cycles. We instantiate this recipe as recAMPLIFY, a pLM with 155M unique parameters trained on the same 2T tokens as AMPLIFY-350M. Despite 56% fewer parameters, it improves CASP contact recovery and three out of four supervised tasks, notably raising fold accuracy by 22.5%. Because optimal depth for recAMPLIFY varies across tasks and assays, we then introduce Gated Mixture Recurrence (GMR), a framework that decodes each protein from a probability-weighted mixture of hidden states visited during recursion. GMR employs a learned residual gate at scalar, low-rank, or per-dimension resolution to control how much of each cycle's proposed update is accepted into the state. This unifies halting and gating into a single adaptive computation objective while recovering PonderNet-style halting as a special case without gates. Layered on top of recAMPLIFY, GMR further reduces expected layer depth by up to 37%, although the accuracy trade-off varies by task. Mechanistic analyses show that residual gating controls recurrence stability beyond the trained depth and that halt depth tracks biological properties.
2026-10-18
arXiv (preprint)
doi.org
arxiv.org
Reconstruction or Semantics? What Makes a Latent Space Useful for Robotic World Models
Nilaksh
Saurav Jha
Artem Zholus
Sarath Chandar
2025-12-31
Neural Information Processing Systems (Accept (poster))
doi.org
openreview.net
Probing the effectiveness of World Models for Spatial Reasoning through Test-time Scaling
Saurav Jha
Muhammad Jehanzeb Mirza
Wei Lin
Shiqi Yang
A. Chandar
Vision-Language Models (VLMs) remain limited in spatial reasoning tasks that require multi-view understanding and embodied perspective shift… (see more)s. Recent approaches such as MindJourney attempt to mitigate this gap through test-time scaling where a world model imagines action-conditioned trajectories and a heuristic verifier selects helpful views from such trajectories. In this work, we systematically examine how such test-time verifiers behave across benchmarks, uncovering both their promise and their pitfalls. Our uncertainty-based analyses show that MindJourney's verifier provides little meaningful calibration, and that random scoring often reduces answer entropy equally well, thus exposing systematic action biases and unreliable reward signals. To mitigate these, we introduce a Verification through Spatial Assertions (ViSA) framework that grounds the test-time reward in verifiable, frame-anchored micro-claims. This principled verifier consistently improves spatial reasoning on the SAT-Real benchmark and corrects trajectory-selection biases through more balanced exploratory behavior. However, on the challenging MMSI-Bench, none of the verifiers, including ours, achieve consistent scaling, suggesting that the current world models form an information bottleneck where imagined views fail to enrich fine-grained reasoning. Together, these findings chart the bad, good, and ugly aspects of test-time verification for world-model-based reasoning. Our code is available at https://github.com/chandar-lab/visa-for-mindjourney.
2025-12-04
ArXiv (preprint)
doi.org
arxiv.org