VertAX: a differentiable vertex model for learning epithelial tissue mechanics

TL;DR

VertAX is a flexible, differentiable framework built with JAX that enables simulation, parameter inference, and inverse design of epithelial tissue mechanics using vertex models.

Contribution

It introduces VertAX, a novel differentiable vertex-modeling framework that integrates automatic differentiation, GPU acceleration, and bilevel optimization for tissue mechanics analysis.

Findings

VertAX successfully models tissue morphogenesis.

It accurately infers mechanical parameters from data.

Equilibrium propagation approximates gradients with limited simulation.

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.