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).
Differentiable simulation is a key component in learning, control, and inverse problems, where gradients through nonlinear implicit solvers are required. Existing approaches either rely on unrolled automatic differentiation, whose memory grows with solver depth, or on equation-level implicit differentiation, which assembles global Jacobians and solves large sparse adjoint systems, discarding the locality of the forward solver -- and differentiating the converged equation rather than the finite computation that actually ran. We propose solver-level differentiation, which differentiates the executed solver itself. When a solver is composed of block implicit updates, its discrete adjoint is obtained by applying the corresponding adjoint updates in reverse order, yielding a reverse-sweep formulation whose backward pass mirrors the forward solver. From an operator perspective, the forward pass realizes an approximate inverse through ordered local solves, and the backward applies its transpose through reverse local adjoint solves, constructing no global system. We instantiate this idea on Vertex Block Descent, yielding a differentiable solver whose reverse colored Gauss-Seidel sweeps are composed entirely of local 3×3 adjoint solves. The backward matches automatic differentiation through the identical executed forward to machine precision at every solver depth, where the equation-level adjoint is off by 37% after one sweep; in a controlled same-codebase, same-GPU comparison it is 33x faster and uses 71x less memory than unrolled automatic differentiation; and the same construction is exact on projective dynamics and extended position-based dynamics. We scale differentiable elastodynamics to 106 contact-coupled soft bodies (8M vertices) on one GPU. Overall, this work highlights solver structure as a practical organizing principle for efficient differentiable simulation.
Lei Shu, Ying Jiang, Kui Wu +3
Stanford University · University of California, Los Angeles · LightSpeed Studios, Tencent America +1
Accurate soft tissue simulation is essential for surgical training, pre-operative planning, and haptic feedback systems. While learning-based surrogate models trained on data using the finite element method (FEM) offer a promising path to real-time inference, their reliability depends on well-calibrated constitutive models. Existing approaches neither provide systematic guidance on model selection across stiffness levels, nor generalize across different tissue stiffnesses or geometries. We perform a comprehensive calibration of hyperelastic constitutive models in the SOFA Framework using gravity-loaded silicone beams with different stiffnesses. Using calibrated simulations as training data, we use a softness conditioned equivariant graph neural network, enabling deformation and force prediction across multiple tissue types and unseen geometries. Our model achieves sub-millimeter mean deformation accuracy at 0.010s inference time, while showing that force prediction quality is directly tied to upstream calibration consistency.
Madina Kojanazarova, Sidaty El Hadramy, Philippe C. Cattin
University of Basel, Department of Biomedical Engineering, Allschwil, Switzerland
Reduced-order models based on Cosserat-rod theory are now well established, and modeling theory is no longer the primary bottleneck in soft-robot control. Their implementations, however, do not support the differentiable, GPU-parallel, and control-oriented workflows that underpin advanced rigid-robotics applications. Here, we fill this gap with SoRoMoX (Soft Robot Models in JAX), a fully numerical, JIT-compilable Python/JAX framework. SoRoMoX implements articulated, Piecewise Constant Strain, and Variable Strain models through a unified, control-ready interface that provides inertia matrices, gravitational and elastic forces, Jacobians, and their derivatives. To our knowledge, it is the first rod/strain-based soft-robot modeling framework that runs directly on GPUs, with Warp kernels accelerating parallel continuum-model execution, and is end-to-end differentiable with respect to states, inputs, and parameters. Sequential CPU rollouts are up to 27.0 times faster than SoRoSim, while GPU-parallel GVS rollouts increase throughput by up to 679.7 times. This performance enables workflows that were previously impractical or impossible: static-equilibrium system identification with 66% lower marker RMSE; residual-force learning with a further 64% reduction; computed-torque tracking with RMSE reduced by a factor of approximately 500 relative to model-free PD; control-gain optimization with up to 98% lower loss than untuned gains; safety-constrained control using high-order control barrier functions to keep the peak contact force within a prescribed 5 N bound, compared with 33.5 N without the safety constraint; and reinforcement-learning policy training up to 7 times faster than a CPU PyElastica discrete-rod baseline through massively parallel rollouts.
Maximilian Stölzle, Solange Gribonval, Daniel Feliu-Talegon +9
Delft University of Technology, Delft, The Netherlands · Italian Institute of Technology, Genoa, Italy · Massachusetts Institute of Technology, Cambridge, MA, USA +1