VertAX: a differentiable vertex model for learning epithelial tissue mechanics
Authors: Alessandro Pasqui, Jim Martin Catacora Ocana, Anshuman Sinha, Matthieu Perez, Fabrice Delbary, Giorgio Gosti, Mattia Miotto, Domenico Caudo, +2 more
Organizations: Center for Interdisciplinary Research in Biology (CIRB), Collège de France, Université PSL, CNRS, INSERM, 75005 Paris, France. · Current address: Center for Integrative Biology (CBI), University of Toulouse, CNRS, Toulouse, France. · Center for Life Nano & Neuro Science, IIT, 00161, Rome, Italy. · Current address: Institute of Heritage Science, CNR, 00010, Montelibretti, Italy. · Department of Physics, Università Sapienza, 00185, Rome, Italy. · Google DeepMind‡.
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.
Figures & tables
Figure 1: VertAX concept and forward/inverse modeling. (a) In VertAX , tissue equilibrium is obtained by minimizing an energy E(X,θ) with respect to vertex positions X ; in inverse problems, model parameters θ are optimized by minimizing a cost C relative to a target. (b) Simulation modes with periodic or curved boundary conditions. (c) Forward modeling: different energy formulations produce different equilibrium geometries from the same initial state. (d) Inverse modeling applications: force inference, inverse mechanical design, and inverse patterning.
Figure 2: Benchmarking bilevel optimization strategies for mechanical parameter inference in VertAX . (a) Synthetic benchmark for inference of target shape factors pα0 . Left, inferred versus ground-truth values at epoch 0; right, inferred versus ground-truth values at epoch 350. Results are shown for AD (blue), EP (green), Adj-ID (red), and Sen-ID (purple). The dashed diagonal indicates perfect agreement. The benchmark uses the energy E1 and the v2v cost. (b, c) Evolution over training epochs of the mean cost function C (b, lower is better) and mean relative error (c, lower is better). (d, e) Box plots of the Pearson correlation coefficient (d, higher is better) and relative error (e, lower is better) at epoch 350. (f, g) Mean runtime per epoch (f, lower is better; log scale) and peak memory usage (g, lower is better; log scale) as a function of cell number. Experimental details: Data in (b–g) show mean ± s.d. over n=5 independent replicates. Unless otherwise specified in (f, g), tissues contained N=20 cells with Lbox=N . Outer-loop optimization used ADAM with learning rate 10−3 , and inner-loop energy minimization used gradient descent with learning rate 10−2 . Inner-loop stopping parameters were a maximum of 103 iterations, tolerance 10−3 , patience 10 , and minimum T1 distance 5×10−3 . For EP, the optimal finite difference parameter β was pre-selected to maximize performance (Extended Figure 1).
Figure 3: Structural and image-based losses for inverse vertex modeling in VertAX . (a) Comparison of two outer-loss functions for bilevel parameter inference from tissue configurations initialized far from the target: a geometric vertex-to-vertex (v2v) cost and an index-aware structural (IAS) cost based on the tissue graph topology. (b, c) Evolution over training epochs of the Pearson correlation coefficient between inferred and ground-truth parameters (b) and of the relative error (c) for inference with the v2v cost alone or with a combined loss (IAS+v2v). (d, e) Box plots of the Pearson correlation coefficient (d) and relative error (e) at epoch 350 across n=5 independent initializations. (f) Schematic of the microscopy reconstruction loss. The polygonal vertex mesh X is convolved with a microscope point-spread function PSF(θC) to generate a synthetic image Isyn(X,θC) . The outer cost C(X,Iobs,θC) is defined from the discrepancy between this synthetic image and the microscopy image Iobs . (g) Application to tissue fluidization inference from experimental images. Left: fraction of zero-energy states as a function of the target shape factor p0 from forward simulations, with sigmoid fit shown in red ( n=100 independent simulations of 100-cell tissues). Center: representative control and TGF β -treated A549 monolayers, with brightfield and E-cadherin/DAPI images, together with the inferred equilibrium mesh initialized from a Voronoi tessellation of Cellpose-segmented nuclei and optimized using the microscopy reconstruction loss. Right: box plots of inferred tissue shape factors for control and TGF β conditions (Mann–Whitney U test, p=0.0071 ).
Figure 4: Inverse design of epithelial tissues with VertAX . (a) Inverse mechanical design of target shape factor using a cost C that penalizes deviations of the mean equilibrium shape factor from Ptarget0=4.10 . Left and center: initial and final equilibrium configurations. Right: epoch-wise evolution of the mean shape factor and of the individual target shape factors pα0 (inset, early-epoch dynamics). (b) Inverse morphological design of tissue elongation using a cost C based on the squared deviation of the aspect ratio r=L/ℓ from a target value rtarget . Left and center: initial ( r=1.50 ) and final ( r=2.75 ) tissue configurations. Right: aspect-ratio mean squared error ( ± std) versus rtarget for tissues of 6, 12, and 18 cells. (c) Inverse patterning design using a cost C=∑edgesℓ11+ℓ22 to favor heterotypic interfaces. Left and center: initial random and final optimized fate patterns. Right: mean ± std of the relative heterotypic interface length ℓ12/ℓtot over epochs ( n=5 replicates).
Delft University of Technology, Delft, The Netherlands · Italian Institute of Technology, Genoa, Italy · Massachusetts Institute of Technology, Cambridge, MA, USA +1