Diffusion with Attention for Inverse Optimization

8 pages•Published: April 19, 2026

Abstract

Traditional methods for inverse optimization of non-linear PDE systems face challenges such as getting trapped in local minima. An AI-based approach can be an alternative to drive optimization through a stable diffusion learning process with a more global context. The concept of using diffusion for generating new control sequences and control functions has already been lightly explored in the literature, however, this architecture has primarily been tested on smooth control functions in PDE simulations. In this paper, we test the limits of attention-based diffusion in inverting a 2D heterogeneous control function coupled in an advection-diffusion PDE system. Recent work has noted that methods such as supervised learning and reinforcement learning have proven somewhat effective; however, they often produce non-physical dynamics or fail to remain optimal long-term. This paper tests a UNet-based diffusion model that uses attention to solve inverse optimization problems. These encouraging results suggest that attention could potentially be an effective mechanism for inverse optimization.

Keyphrases: advection diffusion pde system, ai for physics, diffusion for control, inverse optimization

In: Jernej Masnec, Hamid Reza Karimian, Parisa Kordjamshidi and Yan Li (editors). Proceedings of AI for Accelerated Research Symposium, vol 3, pages 113-120.

Links:	https://easychair.org/publications/paper/LfV1
	https://doi.org/10.29007/twm6