publications

2024

1. ICLR’24 Conference

   Efficient Integrators for Diffusion Generative Models

   Kushagra Pandey, Maja Rudolph , and Stephan Mandt

   International Conference on Learning Representations, 2024

   Abs PDF Code

   Diffusion models suffer from slow sample generation at inference time. Therefore, developing a principled framework for fast deterministic/stochastic sampling for a broader class of diffusion models is a promising direction. We propose two complementary frameworks for accelerating sample generation in pretrained models: Conjugate Integrators and Splitting Integrators. Conjugate integrators generalize DDIM, mapping the reverse diffusion dynamics to a more amenable space for sampling. In contrast, splitting-based integrators, commonly used in molecular dynamics, reduce the numerical simulation error by cleverly alternating between numerical updates involving the data and auxiliary variables. After extensively studying these methods empirically and theoretically, we present a hybrid method that leads to the best-reported performance for diffusion models in augmented spaces. Applied to Phase Space Langevin Diffusion [Pandey & Mandt, 2023] on CIFAR-10, our deterministic and stochastic samplers achieve FID scores of 2.11 and 2.36 in only 100 network function evaluations (NFE) as compared to 2.57 and 2.63 for the best-performing baselines, respectively. Our code and model checkpoints will be made publicly available at https://github.com/mandt-lab/PSLD

2023

1. NeurIPS’23 Workshop

   Towards Fast Stochastic Sampling in Diffusion Generative Models

   Kushagra Pandey, Maja Rudolph , and Stephan Mandt

   NeurIPS 2023 Workshop on Diffusion Models, 2023

   Abs PDF Code

   Diffusion models suffer from slow sample generation at inference time. Despite recent efforts, improving the sampling efficiency of stochastic samplers for diffusion models remains a promising direction. We propose Splitting Integrators for fast stochastic sampling in pre-trained diffusion models in augmented spaces. Commonly used in molecular dynamics, splitting-based integrators attempt to improve sampling efficiency by cleverly alternating between numerical updates involving the data, auxiliary, or noise variables. However, we show that a naive application of splitting integrators is sub-optimal for fast sampling. Consequently, we propose several modifications to improve sampling efficiency and denote the resulting samplers as Reduced Splitting Integrators. In the context of Phase Space Langevin Diffusion [Pandey & Mandt, 2023] on CIFAR-10, our stochastic sampler achieves an FID score of 2.36 in only 100 network function evaluations (NFE) as compared to 2.63 for the best baselines.

2. ICCV’23 Conference

   A Complete Recipe for Diffusion Generative Models

   Kushagra Pandey, and Stephan Mandt

   International Conference on Computer Vision, 2023

   Abs PDF Code

   Score-based Generative Models (SGMs) have demonstrated exceptional synthesis outcomes across various tasks. However, the current design landscape of the forward diffusion process remains largely untapped and often relies on physical heuristics or simplifying assumptions. Utilizing insights from the development of scalable Bayesian posterior samplers, we present a complete recipe for formulating forward processes in SGMs, ensuring convergence to the desired target distribution. Our approach reveals that several existing SGMs can be seen as specific manifestations of our framework. Building upon this method, we introduce Phase Space Langevin Diffusion (PSLD), which relies on score-based modeling within an augmented space enriched by auxiliary variables akin to physical phase space. Empirical results exhibit the superior sample quality and improved speed-quality trade-off of PSLD compared to various competing approaches on established image synthesis benchmarks. Remarkably, PSLD achieves sample quality akin to state-of-the-art SGMs (FID: 2.10 for unconditional CIFAR-10 generation). Lastly, we demonstrate the applicability of PSLD in conditional synthesis using pre-trained score networks, offering an appealing alternative as an SGM backbone for future advancements. Code and model checkpoints can be accessed at https://github.com/mandt-lab/PSLD

2022

1. TMLR Journal

   DiffuseVAE: Efficient, Controllable and High-Fidelity Generation from Low-Dimensional Latents

   Kushagra Pandey, Avideep Mukherjee , Piyush Rai , and 1 more author

   Transactions on Machine Learning Research, 2022

   Abs PDF Code

   Diffusion Probabilistic models have been shown to generate state-of-the-art results on several competitive image synthesis benchmarks but lack a low-dimensional, interpretable latent space, and are slow at generation. On the other hand, Variational Autoencoders (VAEs) typically have access to a low-dimensional latent space but exhibit poor sample quality. Despite recent advances, VAEs usually require high-dimensional hierarchies of the latent codes to generate high-quality samples. We present DiffuseVAE, a novel generative framework that integrates VAE within a diffusion model framework, and leverage this to design a novel conditional parameterization for diffusion models. We show that the resulting model can improve upon the unconditional diffusion model in terms of sampling efficiency while also equipping diffusion models with the low-dimensional VAE inferred latent code. Furthermore, we show that the proposed model can generate high-resolution samples and exhibits synthesis quality comparable to state-of-the-art models on standard benchmarks. Lastly, we show that the proposed method can be used for controllable image synthesis and also exhibits out-of-the-box capabilities for downstream tasks like image super-resolution and denoising. For reproducibility, our source code is publicly available at this link.

2. NAR Journal

   Inference of cell state transitions and cell fate plasticity from single-cell with MARGARET

   Kushagra Pandey, and Hamim Zafar

   Nucleic Acids Research, May 2022

   Abs PDF Code

   Despite recent advances in inferring cellular dynamics using single-cell RNA-seq data, existing trajectory inference (TI) methods face difficulty in accurately reconstructing the cell-state manifold and cell-fate plasticity for complex topologies. Here, we present MARGARET (https://github.com/Zafar-Lab/Margaret) for inferring single-cell trajectory and fate mapping for diverse dynamic cellular processes. MARGARET reconstructs complex trajectory topologies using a deep unsupervised metric learning and a graph-partitioning approach based on a novel connectivity measure, automatically detects terminal cell states, and generalizes the quantification of fate plasticity for complex topologies. On a diverse benchmark consisting of synthetic and real datasets, MARGARET outperformed state-of-the-art methods in recovering global topology and cell pseudotime ordering. For human hematopoiesis, MARGARET accurately identified all major lineages and associated gene expression trends and helped identify transitional progenitors associated with key branching events. For embryoid body differentiation, MARGARET identified novel transitional populations that were validated by bulk sequencing and functionally characterized different precursor populations in the mesoderm lineage. For colon differentiation, MARGARET characterized the lineage for BEST4/OTOP2 cells and the heterogeneity in goblet cell lineage in the colon under normal and inflamed ulcerative colitis conditions. Finally, we demonstrated that MARGARET can scale to large scRNA-seq datasets consisting of ∼ millions of cells.

2021

1. NeurIPS’21 Workshop

   VAEs meet Diffusion Models: Efficient and High-Fidelity Generation

   Kushagra Pandey, Avideep Mukherjee , Piyush Rai , and 1 more author

   NeurIPS’21, Workshop on Deep Generative Models and Downstream Applications, May 2021

   Abs PDF Code

   Diffusion Probabilistic models have been shown to generate state-of-the-art results on several competitive image synthesis benchmarks but lack a low-dimensional, interpretable latent space, and are slow at generation. On the other hand, Variational Autoencoders (VAEs) have access to a low-dimensional latent space but exhibit poor sample quality. Despite recent advances, VAEs usually require large dimensional hierarchies of the latent codes to generate high-quality samples. We present VAEDM, a novel generative framework for \textitrefining VAE generated samples using diffusion models while also presenting a novel conditional forward process parameterization for diffusion models. We show that the resulting parameterization can improve upon the unconditional diffusion model in terms of sampling efficiency during inference while also equipping diffusion models with the low-dimensional VAE inferred latent code.