# VertAX: a differentiable vertex model for learning epithelial tissue mechanics

arXiv: https://arxiv.org/abs/2604.06896 · HTML: https://mlsanity.com/paper/2604.06896/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 70 references that are in the map.
- First posted on arXiv 8 Apr 2026.

## Builds on (most cited first, 17 of 17)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Equilibrium Propagation: Bridging the Gap Between Energy-Based Models and Backpropagation](https://mlsanity.com/paper/1602.05179/) (2016)
- [DiffTaichi: Differentiable Programming for Physical Simulation](https://mlsanity.com/paper/1910.00935/) (2019)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Optimal Transport for structured data with application on graphs](https://mlsanity.com/paper/1805.09114/) (2018)
- [A Review on Bilevel Optimization: From Classical to Evolutionary Approaches and Applications](https://mlsanity.com/paper/1705.06270/) (2017)
- [A Differentiable Programming System to Bridge Machine Learning and Scientific Computing](https://mlsanity.com/paper/1907.07587/) (2019)
- [Scaling Equilibrium Propagation to Deep ConvNets by Drastically Reducing its Gradient Estimator Bias](https://mlsanity.com/paper/2006.03824/) (2020)
- [Energy-based learning algorithms for analog computing: a comparative study](https://mlsanity.com/paper/2312.15103/) (2023)
- [Updates of Equilibrium Prop Match Gradients of Backprop Through Time in an RNN with Static Input](https://mlsanity.com/paper/1905.13633/) (2019)
- [Equivalence of Equilibrium Propagation and Recurrent Backpropagation](https://mlsanity.com/paper/1711.08416/) (2017)
- [Differentiable Physics Models for Real-world Offline Model-based Reinforcement Learning](https://mlsanity.com/paper/2011.01734/) (2020)
- [Beyond backpropagation: bilevel optimization through implicit differentiation and equilibrium propagation](https://mlsanity.com/paper/2205.03076/) (2022)
- [Real-Time Model Calibration with Deep Reinforcement Learning](https://mlsanity.com/paper/2006.04001/) (2020)
- [Engineering morphogenesis of cell clusters with differentiable programming](https://mlsanity.com/paper/2407.06295/) (2024)

## Abstract (arXiv)

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.