← Search

Molei Tao

43 accepted papers

2026

AlignFlow: Improving Flow-based Generative Models with Semi-Discrete Optimal Transport

ICLR 2026poster

Flow-based Generative Models (FGMs) effectively transform noise into a data distribution, and coupling the noise and data in the training of FGM by Optimal Transport (OT) improves the straightness of the flow paths. However, existing OT- based couplings are difficult to combine with modern models an…

Cited by 0SourcecodeScholar
2026

Complexity Analysis of Normalizing Constant Estimation: from Jarzynski Equality to Annealed Importance Sampling and beyond

ICLR 2026poster

Given an unnormalized probability density $\pi\propto\mathrm{e}^{-V}$, estimating its normalizing constant $Z=\int_{\mathbb{R}^d}\mathrm{e}^{-V(x)}\mathrm{d}x$ or free energy $F=-\log Z$ is a crucial problem in Bayesian statistics, statistical mechanics, and machine learning. It is challenging espec…

Cited by 0SourceScholar
2026

Discrete Adjoint Schrödinger Bridge Sampler

ICML 2026poster

Learning discrete neural samplers is challenging due to the lack of gradients and combinatorial complexity. While stochastic optimal control (SOC) and Schrödinger bridge (SB) provide principled solutions, efficient SOC solvers like adjoint matching (AM), which excel in continuous domains, remain une…

Cited by 0SourceScholar
2026

Efficient Diffusion Models under Nonconvex Equality and Inequality constraints via Landing

ICML 2026spotlight

Generative modeling within constrained sets is essential for scientific and engineering applications involving physical, geometric, or safety requirements (e.g., molecular generation, robotics). We present a unified framework for constrained diffusion models on generic nonconvex feasible sets $\Sigm…

Cited by 0SourceScholar
2026

Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models

ICML 2026poster

Inference-time compute has re-emerged as a practical way to improve LLM reasoning. Most test-time scaling (TTS) algorithms rely on autoregressive decoding, which is ill-suited to discrete diffusion language models (dLLMs) due to their parallel decoding over the entire sequence. As a result, developi…

Cited by 0SourceScholar
2026

Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization

ICML 2026spotlight

Diffusion large language models (dLLMs) are promising alternatives to autoregressive large language models (AR-LLMs), as they potentially allow higher inference throughput. Reinforcement learning (RL) is a crucial component for dLLMs to achieve comparable performance with AR-LLMs on important tasks,…

Cited by 0SourceScholar
2026

Generalization of Diffusion Models Arises with a Balanced Representation Space

ICLR 2026poster

Diffusion models generate high-quality, diverse images with great generalizability, yet when overfit to the training objective, they may memorize training samples. We analyze memorization and generalization of diffusion models through the lens of representation learning. Using a two-layer ReLU denoi…

Cited by 0SourcecodeScholar
2026

Improving Classifier-Free Guidance in Masked Diffusion: Low-Dim Theoretical Insights with High-Dim Impact

ICLR 2026poster

Classifier-Free Guidance (CFG) is a widely used technique for conditional generation and improving sample quality in continuous diffusion models, and its extensions to discrete diffusion has recently started to be investigated. In order to improve the algorithms in a principled way, this paper start…

Cited by 0SourceScholar
2026

Improving Reasoning for Diffusion Language Models via Group Diffusion Policy Optimization

ICLR 2026poster

Diffusion language models (DLMs) enable parallel, order-agnostic generation with iterative refinement, offering a flexible alternative to autoregressive large language models (LLMs). However, adapting reinforcement learning (RL) fine-tuning to DLMs remains an open challenge because of the intractabl…

Cited by 0SourceScholar
2026

Lavida-R1: Advancing Reasoning for Unified Multimodal Diffusion Language Models

ICML 2026poster

Diffusion language models (dLLMs) recently emerged as a promising alternative to auto-regressive LLMs. The latest works further extended it to multimodal understanding and generation tasks. In this work, we propose LaViDa-R1, a multimodal, general-purpose reasoning dLLM. Unlike existing works that b…

Cited by 0SourceScholar
2026

MetaDNS: Enhancing Exploration in Discrete Neural Samplers via Metadynamics

ICML 2026poster

Sampling from discrete distributions with multiple modes and energy barriers is fundamental to machine learning and computational physics. Recent discrete neural samplers like MDNS suffer from mode collapse and fail to sample high-energy barrier regions between modes, which is critical for free ener…

Cited by 0SourceScholar
2026

Provable Separations between Memorization and Generalization in Diffusion Models

ICLR 2026poster

Diffusion models have achieved remarkable success across diverse domains, but they remain vulnerable to memorization---reproducing training data rather than generating novel outputs. This not only limits their creative potential but also raises concerns about privacy and safety. While empirical stud…

Cited by 0SourceScholar
2026

Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design

ICML 2026poster

Reinforcement learning has been widely applied to diffusion and flow models for visual tasks such as text-to-image generation. However, these tasks remain challenging because diffusion models have intractable likelihoods, which creates a barrier for directly applying popular policy-gradient type met…

Cited by 0SourceScholar
2025

A Closer Look at Model Collapse: From a Generalization-to-Memorization Perspective

NeurIPS 2025spotlight

The widespread use of diffusion models has led to an abundance of AI-generated data, raising concerns about model collapse---a phenomenon in which recursive iterations of training on synthetic data lead to performance degradation. Prior work primarily characterizes this collapse via variance shrinka…

Cited by 0SourceScholar
2025

Diffuse Everything: Multimodal Diffusion Models on Arbitrary State Spaces

ICML 2025poster

Diffusion models have demonstrated remarkable performance in generating unimodal data across various tasks, including image, video, and text generation. On the contrary, the joint generation of multimodal data through diffusion models is still in the early stages of exploration. Existing approaches…

2025

Diffusion Generative Modeling for Spatially Resolved Gene Expression Inference from Histology Images

ICLR 2025poster

Spatial Transcriptomics (ST) allows a high-resolution measurement of RNA sequence abundance by systematically connecting cell morphology depicted in Hematoxylin and eosin (H\&E) stained histology images to spatially resolved gene expressions. ST is a time-consuming, expensive yet powerful experiment…

2025

Fast Non-Log-Concave Sampling under Nonconvex Equality and Inequality Constraints with Landing

NeurIPS 2025poster

Sampling from constrained statistical distributions is a fundamental task in various fields including Bayesian statistics, computational chemistry, and statistical physics. This article considers the cases where the constrained distribution is described by an unconstrained density, as well as additi…

Cited by 0SourceScholar
2025

Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms

NeurIPS 2025poster

Discrete diffusion models have emerged as a powerful generative modeling framework for discrete data with successful applications spanning from text generation to image synthesis. However, their deployment faces challenges due to the high dimensionality of the state space, necessitating the developm…

Cited by 0SourcecodeScholar
2025

MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control

NeurIPS 2025poster

We study the problem of learning a neural sampler to generate samples from discrete state spaces where the target probability mass function $\pi\propto\mathrm{e}^{-U}$ is known up to a normalizing constant, which is an important task in fields such as statistical physics, machine learning, combinato…

Cited by 0SourcecodeScholar
2025

Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groups

ICLR 2025poster

The generative modeling of data on manifolds is an important task, for which diffusion models in flat spaces typically need nontrivial adaptations. This article demonstrates how a technique called `trivialization' can transfer the effectiveness of diffusion models in Euclidean spaces to Lie groups.…

2025

Variational Learning Finds Flatter Solutions at the Edge of Stability

NeurIPS 2025spotlight

Variational Learning (VL) has recently gained popularity for training deep neural networks. Part of its empirical success can be explained by theories such as PAC-Bayes bounds, minimum description length and marginal likelihood, but little has been done to unravel the implicit regularization in play…

Cited by 0SourceScholar
2024

Extragradient Type Methods for Riemannian Variational Inequality Problems

AISTATS 2024poster

In this work, we consider monotone Riemannian Variational Inequality Problems (RVIPs), which encompass both Riemannian convex optimization and minimax optimization as particular cases. In Euclidean space, the last-iterates of both the extragradient (EG) and past extragradient (PEG) methods converge…

Cited by 7SourcePDFScholar
2024

Provable Acceleration of Nesterov's Accelerated Gradient for Asymmetric Matrix Factorization and Linear Neural Networks

NeurIPS 2024poster

We study the convergence rate of first-order methods for rectangular matrix factorization, which is a canonical nonconvex optimization problem. Specifically, given a rank-$r$ matrix $\mathbf{A}\in\mathbb{R}^{m\times n}$, we prove that gradient descent (GD) can find a pair of $\epsilon$-optimal solut…

Cited by 1SourcePDFScholar
2024

Zeroth-Order Sampling Methods for Non-Log-Concave Distributions: Alleviating Metastability by Denoising Diffusion

NeurIPS 2024poster

This paper considers the problem of sampling from non-logconcave distribution, based on queries of its unnormalized density. It first describes a framework, Denoising Diffusion Monte Carlo (DDMC), based on the simulation of a denoising diffusion process with its score function approximated by a gene…

2023

Deep Momentum Multi-Marginal Schrödinger Bridge

NeurIPS 2023poster

It is a crucial challenge to reconstruct population dynamics using unlabeled samples from distributions at coarse time intervals. Recent approaches such as flow-based models or Schrödinger Bridge (SB) models have demonstrated appealing performance, yet the inferred sample trajectories either fail to…

2023

Mirror Diffusion Models for Constrained and Watermarked Generation

NeurIPS 2023poster

Modern successes of diffusion models in learning complex, high-dimensional data distributions are attributed, in part, to their capability to construct diffusion processes with analytic transition kernels and score functions. The tractability results in a simulation-free framework with stable regres…

2023

Momentum Stiefel Optimizer, with Applications to Suitably-Orthogonal Attention, and Optimal Transport

ICLR 2023poster

The problem of optimization on Stiefel manifold, i.e., minimizing functions of (not necessarily square) matrices that satisfy orthogonality constraints, has been extensively studied. Yet, a new approach is proposed based on, for the first time, an interplay between thoughtfully designed continuous a…

2022

Large Learning Rate Tames Homogeneity: Convergence and Balancing Effect

ICLR 2022poster

Recent empirical advances show that training deep models with large learning rate often improves generalization performance. However, theoretical justifications on the benefits of large learning rate are highly limited, due to challenges in analysis. In this paper, we consider using Gradient Descent…

Cited by 55SourcePDFScholar
2021

Data-driven Prediction of General Hamiltonian Dynamics via Learning Exactly-Symplectic Maps

ICML 2021spotlight

We consider the learning and prediction of nonlinear time series generated by a latent symplectic map. A special case is (not necessarily separable) Hamiltonian systems, whose solution flows give such symplectic maps. For this special case, both generic approaches based on learning the vector field…

Cited by 64SourcePDFScholar
2020

Stochasticity of Deterministic Gradient Descent: Large Learning Rate for Multiscale Objective Function

NeurIPS 2020poster

This article suggests that deterministic Gradient Descent, which does not use any stochastic gradient approximation, can still exhibit stochastic behaviors. In particular, it shows that if the objective function exhibit multiscale behaviors, then in a large learning rate regime which only resolves t…

Cited by 29SourcePDFScholar
2020

Variational Optimization on Lie Groups, with Examples of Leading (Generalized) Eigenvalue Problems

AISTATS 2020poster

The article considers smooth optimization of functions on Lie groups. By generalizing NAG variational principle in vector space (Wibisono et al., 2016) to general Lie groups, continuous Lie-NAG dynamics which are guaranteed to converge to local optimum are obtained. They correspond to momentum versi…

Cited by 26SourcePDFScholar
2020

Why Do Deep Residual Networks Generalize Better than Deep Feedforward Networks? --- A Neural Tangent Kernel Perspective

NeurIPS 2020poster

Deep residual networks (ResNets) have demonstrated better generalization performance than deep feedforward networks (FFNets). However, the theory behind such a phenomenon is still largely unknown. This paper studies this fundamental problem in deep learning from a so-called ``neural tangent kernel''…

Cited by 119SourcePDFScholar