🤖 AI Summary
Traditional vertex models suffer from a large number of parameters, hindering efficient mechanical parameter inference and inverse design. This work proposes the first end-to-end differentiable vertex modeling framework, built on JAX, which integrates automatic differentiation, implicit differentiation, and adjoint-free equilibrium propagation to enable forward simulation, parameter inference, and inverse design of tissue-scale behaviors. The framework allows users to define custom energy functions in pure Python and leverages GPU acceleration for scalable computation. Experimental results demonstrate that the method is both effective and highly scalable across tasks including tissue morphogenesis simulation, mechanical parameter inversion, and inverse design of target structures.
📝 Abstract
Epithelial tissues dynamically reshape through local mechanical interactions among cells, a process well captured by vertex models. Yet their many tunable parameters make inference and optimization challenging, motivating computational frameworks that flexibly model and learn tissue mechanics. We introduce VertAX, a differentiable JAX-based framework for vertex-modeling of confluent epithelia. VertAX provides automatic differentiation, GPU acceleration, and end-to-end bilevel optimization for forward simulation, parameter inference, and inverse mechanical design. Users can define arbitrary energy and cost functions in pure Python, enabling seamless integration with machine-learning pipelines. We demonstrate VertAX on three representative tasks: (i) forward modeling of tissue morphogenesis, (ii) mechanical parameter inference, and (iii) inverse design of tissue-scale behaviors. We benchmark three differentiation strategies-automatic differentiation, implicit differentiation, and equilibrium propagation-showing that the latter can approximate gradients using repeated forward, adjoint-free simulations alone, offering a simple route for extending inverse biophysical problems to non-differentiable simulators with limited additional engineering effort.