Erik Jansson
WASP and KAW postdoctoral fellow in the Cambridge Image Analysis Group and CRA at Wolfson College.
Template deformation methods for reconstruction
Erik Jansson
Idea at a glance
Sinogram (observation)
Reconstruction by deformation
Template
Target
Mathematical language: shape analysis
Other name: matching problems
Provides one possible approach
But what is a shape?
Matching problems: Background
Point cloud
Smooth function
Density
Protein backbones
But what is a shape?
Image from brain scan
Matching problems: Background
A shape is an element of a shape space
A shape space is a space acted on by a suitable Lie group \(G\)
So what's a shape space?
So what's a Lie group?
Matching problems: Background
So what's a Lie group?
Need to know: set of transformations that can be composed and inverted
Examples:
\operatorname{SO}(m): \text{ group of rotations in $m$ dimensions}
\operatorname{Diff}(M): \text{ group of diffeomorphisms of some domain $M$}
\operatorname{GL}(m): \text{ group of invertible linear transformations in $m$ dimensions}
Matching problems: Background
Shapes are in a (metric) space \(V\) acted upon by a Lie group of deformations \(G\).
\min_{\gamma} \underbrace{\mathcal{D}(\gamma.v,w)}_{\text{Matching energy}}+\underbrace{R(\gamma) }_{\text{Regularization}}\\
Find the group element \(\gamma \in G\) minimizing the energy:
Matching problems: Toy example
Goal is to find the optimal group element taking
\(v\)
\(w\)
to (or close to)
Simple example:
G = \operatorname{SO}(2), ~ V = \R^2, \\
v = (1,0), ~ w = (-1/2, 1) \\
\min_{g \in G} \|gv-w\|_{\R^2}^2
\(v\)
\(w\)
\(w\)
\(v\)
\(gv\)
Optimal solution: rotate until aligned
of a shape space
\(v,w\)
\(V\)
Lie Group
\(G\)
acts on elements
Matching problems: Classic setting
Classic case: \(G = \operatorname{Diff}(M)\), \(V = C^\infty(M)\)
All spaces involved are infinite-dimensional
Consider \(\operatorname{Diff}(M)\) as the endpoint of a curve of diffeomorphisms generated by a vector field. Expensive, but geometrically nice (LDDMM!).
\min_{v(t)} \underbrace{\mathcal{D}(\gamma(1).A,B)}_{\text{Matching energy}}+\underbrace{\int_0^1 \int_M L\nu(t) \cdot \nu(t)\mathrm{d}x\mathrm{d}t}_{\text{Regularization}}\\
\dot \gamma = \nu \circ \gamma^{-1}
Matching problems: More interesting case
By calculus of variations magic, the problem reduces down to solving and analysing a tricky PDE
Issue: Slow, as we need to solve two DEs for each gradient step
However, my goal is not matching two shapes...
Indirect matching problems
Reconstruct the object that gave rise to the data.
Lie Group
of a shape space
acts on an element
\(G\)
\(v\)
\(V\)
\(v\)
Match
in a space
by a forward model
with data
\(d\)
\(Y\)
\(\mathcal{F}\colon V \to Y\)
Example from Lang et al, 2018
Sinogram (observation)
Reconstruction by deformation
Target
Template
Reconstruction by classical methods
Issue: slow with LDDMM-approach
Essentially: Solve a variational problem on \(G\), by introducing extra variables,
Iterative updates of infinitesimal generators
One update: solve several DE and compute gradients through these updates
What if we could gradient flow directly on group?
\(\text{iterate}\)
Towards gradient flows
How to compute gradients?
Generally tricky in non-Euclidean setting!
To compute gradients, we need to have a metric.
A clever choice of metric makes the problem tractable.
(\dot g_1, \dot g_2)_g = \langle \mathbb{A} (\dot g_1 \circ g^{-1}), \dot g_2 \circ g^{-1}\rangle_{\mathfrak{g}^*,\mathfrak{g}}, \\
g_i\in T_g G, g_i \in T_e G = \mathfrak{g}\\
\mathbb{A}\colon \mathfrak{g} \to \mathfrak{g}^*
Translate vectors back to algebra, map one to dual with inertia operator, take duality pairing
Gradient computation via momentum map
The Lie group action on the shape space induces a momentum map
\langle J(w,p),u \rangle_{\mathfrak{g}^*,\mathfrak{g}} = \langle p, u_\Phi(w)\rangle_{T^*_w V,T_w V}
Let \(f\colon V \to \R\) and take \(E(g) = f(\Phi(g,w))\)
~\dot g = -\nabla E(g) = -TR_g \mathbb{A}^{-1} J(\Phi(g,w),df(\Phi(g,w))~~
\(\Phi\colon G \times V \to V\) is the group action
Why this metric?
Moral of the story: with this choice of metric, the gradient flow has a closed form expression
Gradient flow for reconstruction
f(w) = \mathcal{L}_y(\mathcal{F}(w)), ~ E(g) = \mathcal{L}_y(\mathcal{F}(\Phi(g,w)))
~\dot g = -TR_g \mathbb{A}^{-1} J(\Phi(g,w),d(\mathcal{L}_y \circ \mathcal{F})(\Phi(g,w))~~
- A function \(f \colon V \to \R\)
- Fix a template \(w \in V\) and construct \(E \colon G \to \mathbb{R}\) by \(E(g) = f(\Phi(g,w))\)
Ingredients:
How to construct \(f\)? Take observed data \(y \in Y\) and some data loss \(\mathcal{L}_y \colon Y \to \mathbb{R}\)
Multiple reasons for studying this
Potentially useful
Nice mathematics
Infinite-dimensional geometry
Nonlinear PDE theory
Well-posedness analysis
Geometric mechanics
Nice way of incorporating
anatomical priors
into inverse problems
Application: Reconstructing from Cryo-EM data
in 30 seconds (by a maths person)
Electron
microscope
Low dose: low SNR
Many images
Many copies of the same protein but at different orientations and conformations
Particle picking
Protein modelling
Proteins in 30 seconds (by a mathematician)
In this work: forget about everything but the \(C_\alpha\)s
Relative positions
\((\mathbb R^3)^N\)
Mathematical model for proteins
Set-up and forward model
In the end: Each image with one projection and forward model \(\mathcal{F}_j, j = 1,\ldots, M\)
Shape space: space of relative positions \(V = \mathbb{R}^{3N}\)
Data space - M 2D images, \(L^2(\mathbb{R}^2)^M\)
We want rigid deformations, so \(G = \operatorname{SO}(3)^N\)
Gradient flow
\dot g = -\eta g \\
\mathbb{A}_i\eta_i=\Phi(g, w)_i\times\left(\sum_j(d \mathcal{F}_j(\Phi(g, w)))^*\left(-2\frac{a_{j}}{b_{j} c_{j}}\left(y_j-a_{j}\frac{\mathcal{F}_j(\Phi(g, w))}{b_{j}}\right)\right) \right)_i, \quad i = 1, \ldots N\\
a_{j} = (\mathcal{F}_j(\Phi(g, w)),y_j)_{L^2}\\
b_{j} = \|y_j\|_{L^2}\\
c_{j} = \left \|\mathcal{F}_j(\Phi(g, w)) \right\|_{L^2}
\(y_1, \ldots y_M\) are images,
\(\Phi\) is the group action of \(g\)
The adjoint is easy to derive but results in rather unaesthetic expression
For each image, use NCC similarity, reconstruct by deforming template \(w\)
Gradient flow
Data and goal
Faint outline of protein can be seen, but \(SNR = 0.01\)
Goal to reconstruct
Template
g \in \operatorname{SO}(3)^N
Find \(g\)
Full protein + 4 Gaussian forward model is used for data generation
3D reconstruction results
2D projection results
Multi-chain
Extension to multi-chain setting is straightforward
Here: 3 Chains, each of ~340 residues, in total 1039 residues
Note on well-posedness
In finite dimensions
(on a compact group)....
Global well-posedness is easy to prove!
As are the properties of the IP
Still: Toy problems
1039 residues is not that much.
Example: EMPIAR-10028 has 80 chains and >10k residues
- Keep improving on the gradient flow method and incorporate orientation estimation.
- Retain the template-deformation based method, but instead use a neural network to deform the template - speedup, but maybe at other costs? (see e.g. arXiv:2602.21915v1)
Template quality?
Infinite-dimensional application: tomography
In this case \(G\) is the set of all diffeomorphisms of the domain, which is inf-dim.
Domain is 2D, so \(\mathbb{R}^2\)
Observations
Ray transform
Angle in [0,60], six angles
Target
Tomography
in 30 seconds (by a maths person)
X-ray tube
Test object
Detector
Rotation
Infinite-dimensional set-up
- Classical set-up: \(G = \operatorname{Diff}(M), V = C^\infty(M), Y =L^2(S^1\times \mathbb{R}) \)
- Forward operator \(\mathcal{F}\colon \to L^2(S^1 \times \mathbb{R})\) is the ray transform.
For one beam parametrized by \((\theta \in S^1, t \in \mathbb{R})\), its value is given by
- ​Data fidelity is once again NCC \(\mathcal{L}_y(y')=1-\frac{\langle y,y'\rangle}{\|y\| \|y'\|}\)
\mathcal{R}_{\theta,t}(\rho) = \int_{\mathbb{R}} \rho(t\omega^\perp + u \omega) \mathrm{d}u, ~\omega = (\cos \theta, \sin \theta)
Momentum map of action \(\Phi(g,w) = w\circ g^{-1}\)
J(I,P) = -P \nabla I
Infinite-dimensional gradient flow
Tomography gradient flow
The unregularized gradient flow is given by:
\dot g = \eta \circ g,\\
\mathbb{A} \eta = d\mathcal{F}(w')^* \left[2\frac{\langle \mathcal{F}(w'),y \rangle }{\|\mathcal{F}(w')\|^2\|y\|^2}\left(\frac{\langle \mathcal{F}(w'),y \rangle\mathcal{F}(w')}{\|\mathcal{F}(w')\|^2}-y\right)\right] \nabla w' \\
w' = \Phi(w,g)
How and why to regularize?
Tomography results
Observations
Reconstruction
Target
Quick!
Tomography results
But now turn up the noise... (a factor 100, so not super realistic)
Reconstruction quality starts to deterioriate!
Framework can handle regularisation as well! But this is a story for another time
Two approaches to regularization
- Set \(V_R = \operatorname{Dens}(M)\): Regularization by action of \(g\) on reference volume, \(\Phi_R(g, \mathrm{d}x) = |Dg^{-1}|\mathrm{d}x\)
d_{\operatorname{FR}}^2(|Dg^{-1}|\mathrm{d}x,\mathrm{d}x) = |M| \arccos^2\left(\frac{1}{|M|} \int_{M} \sqrt{|Dg^{-1}|} \mathrm{d}x\right),
Recovers regularization part of a gradient flow from Bauer et al.
- Set \(V_R = \operatorname{Met}(M)\): Regularization by action of \(g\) on reference metric \(\mathbf{h}\), \(\Phi_R(g, \mathbf{h}) =\mathbf{h}_{g^{-1}(x)} (dg^{-1}u,dg^{-1}v)\)
Recovers regularization part of a gradient flow from Balehowsky et al.
\|\mathbf{h}_1-\mathbf{h}_2\|^2 = \int_{M}((\mathbf{h}_1)^{ij} -(\mathbf{h}_2)^{ij}) [(\mathbf{h}_1)_{ij} -(\mathbf{h}_2)_{ij} ]
Two approaches to regularization
Extend shape space and re-apply momentum map (which now is a sum)
Big question for me: Well-posedness.
For example, in tomography, pseudodifferential operators inside the GF appears.
Bit messy, but maybe Ebin–Marsden saves the day?
V' = V \times V_R\\
J(w_1,p_1,w_2,p_2)' = J(w_1,p_1) + J_R(w_2,p_2)
In summary
Anatomical priors can be incorporated
using template deformation
Gradient flows speeds up reconstruction
compared to LDDMM-based approaches
A general framework: change problem,
change the ingredients, but geometry
stays the same
Diamond light source
By Erik Jansson
