Generative Inverse Design with Abstention via Diagonal Flow Matching

Inverse design aims to find design parameters $x$ achieving target performance $y^*$. Generative approaches learn bidirectional mappings between designs and labels, enabling diverse solution sampling. However, standard conditional flow matching (CFM), when adapted to inverse problems by pairing labels with design parameters, exhibits strong sensitivity to their arbitrary ordering and scaling, leading to unstable training. We introduce Diagonal Flow Matching (Diag--CFM), which resolves this through a zero-anchoring strategy that pairs design coordinates with noise and labels with zero, making the learning problem provably invariant to coordinate permutations. This yields substantially lower round-trip error than CFM and invertible neural network baselines across design dimensions up to $P{=}784$, including order-of-magnitude gains on several benchmarks. We develop two architecture-intrinsic uncertainty metrics, Zero-Deviation and Self-Consistency, that enable three practical capabilities: selecting the best candidate among multiple generations, abstaining from unreliable predictions, and detecting out-of-distribution targets; consistently outperforming ensemble and general-purpose alternatives across all tasks. We validate on airfoil, gas turbine combustor, scalable analytical benchmarks, a photonics inverse-design task, and an image-statistics benchmark.

Engineers often need to design objects by working backwards: instead of testing one proposed design, they start with desired performance and ask what designs could achieve it. Examples include airfoils, gas turbines, optical coatings, and even image-like design spaces. Modern generative AI can propose many candidate designs, but it will usually produce an answer even when the requested target is unrealistic or outside what it learned from data. Our work introduces Diagonal Flow Matching, a way to train such models so that design variables and performance targets are not accidentally tied to arbitrary coordinate orderings. This makes the model more stable and accurate when generating designs from requested performance goals. We also add two built-in warning signals that estimate when a generated design should not be trusted, without requiring a large collection of separately trained models. These signals help choose the best candidate among several generated designs, reject unreliable outputs, and detect requests that are likely outside the feasible design space. Across engineering and benchmark tasks, the method produces more accurate designs and gives users a practical way to abstain instead of silently returning a poor design.