Portrait of Kusha Sareen

Kusha Sareen

Master's Research - McGill University
Supervisor
Siamak Ravanbakhsh
Research Topics
Deep Learning
Machine Learning Theory
Website
GitHub

Publications

CUBE: A Standard for Unifying Agent Benchmarks
Alexandre Lacoste
Nicolas Gontier
Oleh Shliazhko
Aman Jaiswal
Kusha Sareen
Shailesh Nanisetty
Joan Cabezas
Manuel Del Verme
Omar G. Younis
Simone Baratta
Matteo Avalle
Imene Kerboua
Xing Han Lu
Elron Bandel
Michal Shmueli-Scheuer
Asaf Yehudai
Leshem Choshen
Jonathan Lebensold
Sean Hughes
Massimo Caccia … (see 6 more)
Alexandre Drouin
Siva Reddy
Tao Yu
Yu Su
Graham Neubig
Dawn Song
The proliferation of agent benchmarks has created critical fragmentation that threatens research productivity. Each new benchmark requires s… (see more)ubstantial custom integration, creating an "integration tax" that limits comprehensive evaluation. We propose CUBE (Common Unified Benchmark Environments), a universal protocol standard built on MCP and Gym that allows benchmarks to be wrapped once and used everywhere. By separating task, benchmark, package, and registry concerns into distinct API layers, CUBE enables any compliant platform to access any compliant benchmark for evaluation, RL training, or data generation without custom integration. We call on the community to contribute to the development of this standard before platform-specific implementations deepen fragmentation as benchmark production accelerates through 2026.
2026-03-15
arXiv (preprint)
doi.org
arxiv.org
Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models
Vineet Jain
Kusha Sareen
Mohammad Pedramfar
Siamak Ravanbakhsh
Adapting a pretrained diffusion model to new objectives at inference time remains an open problem in generative modeling. Existing steering … (see more)methods suffer from inaccurate value estimation, especially at high noise levels, which biases guidance. Moreover, information from past runs is not reused to improve sample quality, resulting in inefficient use of compute. Inspired by the success of Monte Carlo Tree Search, we address these limitations by casting inference-time alignment as a search problem that reuses past computations. We introduce a tree-based approach that samples from the reward-aligned target density by propagating terminal rewards back through the diffusion chain and iteratively refining value estimates with each additional generation. Our proposed method, Diffusion Tree Sampling (DTS), produces asymptotically exact samples from the target distribution in the limit of infinite rollouts, and its greedy variant, Diffusion Tree Search (DTS
2025-09-17
NeurIPS.cc/2025/Conference (poster)
doi.org
openreview.net
Energy Loss Functions for Physical Systems
Sékou-Oumar Kaba
Kusha Sareen
Daniel Levy
Siamak Ravanbakhsh
Effectively leveraging prior knowledge of a system's physics is crucial for applications of machine learning to scientific domains. Previous… (see more) approaches mostly focused on incorporating physical insights at the architectural level. In this paper, we propose a framework to leverage physical information directly into the loss function for prediction and generative modeling tasks on systems like molecules and spins. We derive energy loss functions assuming that each data sample is in thermal equilibrium with respect to an approximate energy landscape. By using the reverse KL divergence with a Boltzmann distribution around the data, we obtain the loss as an energy difference between the data and the model predictions. This perspective also recasts traditional objectives like MSE as energy-based, but with a physically meaningless energy. In contrast, our formulation yields physically grounded loss functions with gradients that better align with valid configurations, while being architecture-agnostic and computationally efficient. The energy loss functions also inherently respect physical symmetries. We demonstrate our approach on molecular generation and spin ground-state prediction and report significant improvements over baselines.
2025-09-17
NeurIPS.cc/2025/Conference (poster)
doi.org
openreview.net
Putting the Value Back in RL: Better Test-Time Scaling by Unifying LLM Reasoners With Verifiers
Kusha Sareen
Morgane M Moss
Alessandro Sordoni
Rishabh Agarwal
Arian Hosseini
2025-07-06
colmweb.org/COLM/2025/Conference (accepted)
doi.org
openreview.net
Symmetry-Aware Generative Modeling through Learned Canonicalization
Kusha Sareen
Daniel Levy
Arnab Kumar Mondal
Sékou-Oumar Kaba
Tara Akhound-Sadegh
Siamak Ravanbakhsh
Generative modeling of symmetric densities has a range of applications in AI for science, from drug discovery to physics simulations. The ex… (see more)isting generative modeling paradigm for invariant densities combines an invariant prior with an equivariant generative process. However, we observe that this technique is not necessary and has several drawbacks resulting from the limitations of equivariant networks. Instead, we propose to model a learned slice of the density so that only one representative element per orbit is learned. To accomplish this, we learn a group-equivariant canonicalization network that maps training samples to a canonical pose and train a non-equivariant generative model over these canonicalized samples. We implement this idea in the context of diffusion models. Our preliminary experimental results on molecular modeling are promising, demonstrating improved sample quality and faster inference time.
2024-10-22
NeurIPS.cc/2024/Workshop/NeurReps (poster)
doi.org
openreview.net