Leveraging automatic differentiation in modern machine learning frameworks for (neural) topology optimization

Abstract

Automatic differentiation (AD) was introduced into topology optimization (TO) more than two decades ago to compute accurate gradients through complex computational workflows. Nevertheless, its adoption within the TO community has remained limited, largely due to the strong reliance on adjoint-based sensitivity analysis—which typically offers superior memory efficiency and runtime performance—and the practical difficulties of integrating large-scale simulations into specialized AD frameworks. The recent rise of machine learning (ML) has opened new opportunities for TO through the advanced AD capabilities of modern ML frameworks such as JAX and PyTorch. A growing body of work at the intersection of ML and TO now focuses on tightly coupling ML components with classical TO workflows. Neural TO is a prominent example, in which an untrained neural network parameterizes the material density field and optimization proceeds over the network parameters. To enable such ML–TO hybrid workflows, a deeper understanding of how AD systems operate in these frameworks is essential. This article explains the practical principles of AD in modern ML frameworks and their relation to classical adjoint-based sensitivity analysis. We present implementation strategies for wrapping essential operations—such as finite element solvers—into AD-compatible components without reimplementing them from scratch. These ideas are illustrated through two compact code examples: a classical TO pipeline with selectively AD-wrapped components and a neural TO workflow.