Portrait de Kusha Sareen

Kusha Sareen

Maîtrise recherche - McGill
Superviseur⋅e principal⋅e
Siamak Ravanbakhsh
Sujets de recherche
Apprentissage profond
Théorie de l'apprentissage automatique
Site web
GitHub

Publications

Putting the Value Back in RL: Better Test-Time Scaling by Unifying LLM Reasoners With Verifiers
Kusha Sareen
Morgane M Moss
Alessandro Sordoni
Arian Hosseini
Rishabh Agarwal
2025-07-07
colmweb.org/COLM/2025/Conference (accepté)
doi.org
openreview.net
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 … (voir plus)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, leading to 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*) performs a robust search for high reward samples. On MNIST and CIFAR-10 class-conditional generation, DTS matches the FID of the best-performing baseline with up to
2025-06-10
ICML.cc/2025/Workshop/PUT (poster)
openreview.net
openreview.net
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 … (voir plus)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-06-10
ICML.cc/2025/Workshop/PUT (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-05-07
ArXiv (prépublication)
arxiv.org
arxiv.org
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… (voir plus)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.
2025-01-14
ArXiv (prépublication)
arxiv.org
arxiv.org
Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models
Vineet Jain
Kusha Sareen
Mohammad Pedramfar
Siamak Ravanbakhsh
2025-01-01
arXiv.org (prépublication)
doi.org
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… (voir plus)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-23
NeurIPS.cc/2024/Workshop/NeurReps (poster)
doi.org
openreview.net