Publications

2026

1. preprint

   

   Variational Flow Maps: Make Some Noise for One-Step Conditional Generation

   Abbas Mammadov^*, So Takao^*, Bohan Chen , and 4 more authors

   preprint, Mar 2026

   Abstract arXiv Code

   Flow maps enable high-quality image generation in a single forward pass. However, unlike iterative diffusion models, their lack of an explicit sampling trajectory impedes incorporating external constraints for conditional generation and solving inverse problems. We put forth Variational Flow Maps, a framework for conditional sampling that shifts the perspective of conditioning from "guiding a sampling path", to that of "learning the proper initial noise". Specifically, given an observation, we seek to learn a noise adapter model that outputs a noise distribution, so that after mapping to the data space via flow map, the samples respect the observation and data prior. To this end, we develop a principled variational objective that jointly trains the noise adapter and the flow map, improving noise-data alignment, such that sampling from complex data posterior is achieved with a simple adapter. Experiments on various inverse problems show that VFMs produce well-calibrated conditional samples in a single (or few) steps. For ImageNet, VFM attains competitive fidelity while accelerating the sampling by orders of magnitude compared to alternative iterative diffusion/flow models. Code is available at https://github.com/abbasmammadov/VFM.

2. preprint

   

   Meta Flow Maps enable scalable reward alignment

   Peter Potaptchik^*, Adhi Saravanan^*, Abbas Mammadov , and 3 more authors

   preprint, Jan 2026

   Abstract arXiv Code

   Controlling generative models is computationally expensive. This is because optimal alignment with a reward function–whether via inference-time steering or fine-tuning–requires estimating the value function. This task demands access to the conditional posterior p1|t(x1|xt), the distribution of clean data x1 consistent with an intermediate state xt, a requirement that typically compels methods to resort to costly trajectory simulations. To address this bottleneck, we introduce Meta Flow Maps (MFMs), a framework extending consistency models and flow maps into the stochastic regime. MFMs are trained to perform stochastic one-step posterior sampling, generating arbitrarily many i.i.d. draws of clean data x1 from any intermediate state. Crucially, these samples provide a differentiable reparametrization that unlocks efficient value function estimation. We leverage this capability to solve bottlenecks in both paradigms: enabling inference-time steering without inner rollouts, and facilitating unbiased, off-policy fine-tuning to general rewards. Empirically, our single-particle steered-MFM sampler outperforms a Best-of-1000 baseline on ImageNet across multiple rewards at a fraction of the compute.

2025

1. NeurIPS 2025

   

   Guided Diffusion Sampling on Function Spaces with Applications to PDEs

   Jiachen Yao^*, Abbas Mammadov^*, Julius Berner , and 4 more authors

   Neural Information Processing Systems (NeurIPS), May 2025

   Abstract arXiv Code

   We propose a general framework for conditional sampling in PDE-based inverse problems, targeting the recovery of whole solutions from extremely sparse or noisy measurements. This is accomplished by a function-space diffusion model and plug-and-play guidance for conditioning. Our method first trains an unconditional discretization-agnostic denoising model using neural operator architectures. At inference, we refine the samples to satisfy sparse observation data via a gradientbased guidance mechanism. Through rigorous mathematical analysis, we extend Tweedie’s formula to infinite-dimensional Banach spaces, providing the theoretical foundation for our posterior sampling approach. Our method (FunDPS) accurately captures posterior distribution in function spaces under minimal supervision and severe data scarcity. Across five PDE tasks with only 3% observation, our method achieves an average 32% accuracy improvement over state-of-the-art fixed-resolution diffusion baselines while reducing sampling steps by 4x. Furthermore, multi-resolution fine-tuning ensures strong cross-resolution generalizability. To the best of our knowledge, this is the first diffusion-based framework to operate independently of discretization, offering a practical and flexible solution for forward and inverse problems in the context of PDEs. Code is available at https://github.com/neuraloperator/FunDPS.

2. ICMLW 2025

   

   EquiReg: Equivariance Regularized Diffusion for Inverse Problems

   Bahareh Tolooshams^*, Aditi Chandrashekar^*, Rayhan Zirvi^* , and 4 more authors

   Building Physically Plausible World Models Workshop at ICML, May 2025

   Abstract arXiv

   Diffusion models represent the state-of-the-art for solving inverse problems such as image restoration tasks. In the Bayesian framework, diffusion-based inverse solvers incorporate a likelihood term to guide the prior sampling process, generating data consistent with the posterior distribution. However, due to the intractability of the likelihood term, many current methods rely on isotropic Gaussian approximations, which lead to deviations from the data manifold and result in inconsistent, unstable reconstructions. We propose Equivariance Regularized (EquiReg) diffusion, a general framework for regularizing posterior sampling in diffusion-based inverse problem solvers. EquiReg enhances reconstructions by reweighting diffusion trajectories and penalizing those that deviate from the data manifold. We define a new distribution-dependent equivariance error, empirically identify functions that exhibit low error for on-manifold samples and higher error for off-manifold samples, and leverage these functions to regularize the diffusion sampling process. When applied to a variety of solvers, EquiReg outperforms state-of-the-art diffusion models in both linear and nonlinear image restoration tasks, as well as in reconstructing partial differential equations.

3. ICLRW 2025

   

   Amortized Posterior Sampling with Diffusion Prior Distillation

   Abbas Mammadov^*, Hyungjin Chung^*, and Jong Chul Ye

   The Frontiers in Probabilistic Inference: Sampling meets Learning (FPI) workshop at ICLR, Apr 2025

   Abstract arXiv

   We propose a variational inference approach to sample from the posterior distribution for solving inverse problems. From a pre-trained diffusion model, our approach trains a conditional flow model to minimize the divergence between the proposal variational distribution and the posterior distribution implicitly defined through the diffusion model. Once trained, the flow model is capable of sampling from the posterior distribution with a single NFE, amortized with respect to the measurement. The proposed method paves a new path for distilling a diffusion prior for efficient posterior sampling. We show that our method is applicable to standard signals in Euclidean space, as well as signals on manifold.

2024

1. NeurIPSW 2024

   

   Diffusion-Based Inverse Solver on Function Spaces With Applications to PDEs

   Abbas Mammadov, Julius Berner, Kamyar Azizzadenesheli , and 2 more authors

   Machine Learning and the Physical Sciences Workshop at NeurIPS, 2024

   Abstract PDF Poster

   We present a novel framework for solving inverse problems in function spaces using diffusion-based generative models. Unlike traditional methods, which often discretize the domain and operate on fixed grids, our approach is discretization-agnostic, allowing for flexibility during sampling and generalization across different resolutions. Built upon function space diffusion models with neural operator architectures, we adapt the denoising process of pre-trained diffusion models to efficiently sample from posterior distributions in function spaces. This framework can be applied to a variety of problems, such as recovery of initial conditions and coefficient functions in noisy or partially observed PDE-based inverse problems like Darcy flows and Navier–Stokes equations. To the best of our knowledge, this is the first diffusion-based plug-and-play solver for inverse problems that operates in a discretization-agnostic manner, providing a new perspective on inverse problems with functional data, as typically arising in the context of PDEs.

2. preprint

   

   Geometric Diffusion Models for Data Over Arbitrary Manifolds

   Byeongsu Sim^*, Abbas Mammadov^*, Moo K. Chung , and 2 more authors

   preprint, 2024
3. ICML 2024

   

   Defining Neural Network Architecture through Polytope Structures of Dataset

   Sangmin Lee, Abbas Mammadov, and Jong Chul Ye

   International Conference on Machine Learning (ICML) SPOTLIGHT, May 2024

   SPOTLIGHT Abstract arXiv Code

   Top 3.5%

   Current theoretical and empirical research in neural networks suggests that complex datasets require large network architectures for thorough classification, yet the precise nature of this relationship remains unclear. This paper tackles this issue by defining upper and lower bounds for neural network widths, which are informed by the polytope structure of the dataset in question. We also delve into the application of these principles to simplicial complexes and specific manifold shapes, explaining how the requirement for network width varies in accordance with the geometric complexity of the dataset. Moreover, we develop an algorithm to investigate a converse situation where the polytope structure of a dataset can be inferred from its corresponding trained neural networks. Through our algorithm, it is established that popular datasets such as MNIST, Fashion-MNIST, and CIFAR10 can be efficiently encapsulated using no more than two polytopes with a small number of faces.

2023

1. KSC 2023

   

   Artificial Barber: Hair Color and Style transfer using GANs

   Abbas Mammadov, and Kaleb Mesfin Asfaw

   Korea Software Congress (KSC) ORAL, Nov 2023

   Abstract PDF Code

   Despite the recent success of face image generation and conditional hair editing with GANs, there are several limitationsdue to its complex geometry and appearance. These models are quite heavy which leads to late inference time. In addition,those models are only considering to change color to some existing hair color but not dealing with customized color. Tocontinue, taking appropriate image for inference (correct angle and alignment) is also cumbersome. Lastly, those models arequite complex to figure out and use due to absence of convenient user interface. In this paper, we present Artificial Barber,a novel conditional and lightweight machine learning pipeline deployed on Qt based Graphical User Interface to change haircolor and style in a fast, accurate, and convenient manner. All the functionalities can be controlled through the GUI. The code is available at https://github.com/abbasmammadov/Artificial-Barber.