Erik Jansson
WASP and KAW postdoctoral fellow in the Cambridge Image Analysis Group and CRA at Wolfson College.
Gradient flows for
indirect matching
problems
Erik Jansson
Joint work with:
Klas Modin (Chalmers Uni)
Carola-Bibiane Schönlieb
(Cambridge)
Ozan Öktem (KTH)
Jonathan Krook (KTH)
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
Reconstructing from Cryo-EM data
Reconstructing from Cryo-EM data
Issue: SNR is around 0.01-0.001, so almost no information
Could the problem be made easier if we knew which protein
was imaged and that we are observing a deformation of a representative conformation?
Hint: Yes, this incorporates an
"anatomical prior"
Reconstructing from Cryo-EM data
Goal to reconstruct
Template
Reconstructing from Cryo-EM data
3D reconstruction results
2D projection results
Tomography
in 30 seconds (by a maths person)
X-ray tube
Test object
Detector
Rotation
Tomography
We consider the limited angle (angle X-ray tube can slide to most 60 degrees)
sparse (only six angles) tomography.
Data
Target
Two classic reconstructions
"Our" reconstruction
Same principle
We are imaging a brain: anatomical prior
Data
Target
Template
Same principle
Deform a template to match an indirectly observed target
Stage is set for some shape matching!
Shapes are in a (metric) space \(V\) acted upon by a Lie group of deformations \(G\).
\min_{\gamma} \underbrace{\mathcal{D}(\gamma.A,B)}_{\text{Matching energy}}+\underbrace{\lambda d_G(\gamma,e) }_{\text{Regularization}}\\
Quick background (LDDMM)
Deformations \(\gamma\) is endpoint of curve \(\gamma \colon [0,1] \to G \)
\dot \gamma_\nu(t) = \mathrm d R_{\gamma_\nu(t)} {\color{black}\nu(t)},\quad \gamma_\nu(0) = e
\min_{\gamma} \left(d^2_V(\gamma_\nu(1) \cdot A, B)+\lambda \int_0^1 \langle \dot \gamma_\nu(t) ,\dot \gamma_\nu(t)\rangle_{\gamma(t)} \mathrm d t\right)
Right-invariance of metric
\min_{\nu} \left(d^2_V(\gamma_\nu(1) \cdot A, B)+\lambda \int_0^1 \langle \nu(t) , \nu(t)\rangle_{e} \mathrm{d}t\right)
Indirect matching (LDDMM)
Geometric framework remains the same.
- Group \(G\) acting on shape space \(V\)
- Data is in data space \(Y\) (Hilbert)
Include a forward operator \(\mathcal{F}\colon V \to Y\)
\underbrace{d_Y^2(\mathcal{F}(\gamma(1).A),B)}_{\text{Matching energy \\changes}}+\underbrace{\int_0^1 \langle \nu(t),\nu(t) \rangle_e \mathrm{d}t}_{\text{Regularization is the same}}
Reconstructs the \(\gamma(1).A\) that best would map to \(B\)
Indirect matching (LDDMM)
- LDDMM framework remains the same, solve by shooting or your favourite algorithm
- However, framework remains slow (additional variables introduced)
- Gradient contains adjoint of differential of forward map: nonlinear FM => optimization issues
- Experience tell us LDDMM-based reconstruction works very well, though (at least for the two examples seen earlier)
Following Balehoskwy, Karlsson and Modin, there might be a faster way
Gradient flow directly on the group
Indirect matching (LDDMM)
Gradient flow directly on the group
Gradients are Riemannian object, so equip \(G\) with a right-invariant metric.
(\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
Let \(y \in Y\) be observed data
Set \(f = \mathcal{L}_y \circ \mathcal{F} \) for data loss function \(\mathcal{L}_y \colon Y \to \mathbb{R} \) and forward model \(\mathcal{F}\colon V \to Y\)
Then \(E(g) = \mathcal{L} \circ \mathcal{F} \circ \Phi(g,w)\) is a indirect matching loss!
Gradient flows!
We would like to minimize:
E(g) = \mathcal{L}_y(\mathcal{F}(\Phi(g,w)))
So we run:
~\dot g = -TR_g \mathbb{A}^{-1} J(\Phi(g,w),d(\mathcal{L}_y \circ \mathcal{F})(\Phi(g,w))~~
A question: isn't this just greedy matching?
No, because we can easily regularize: Extend shape space \(V' = V \times V_R\), where the action of \(G\) on \(V_R\) is something that we should keep small (i.e., metric, background density, etc. Depends on group and problem)
Back to where we started, proteins
Goal to reconstruct
Template
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
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\)
Action of rotations on relative positions ensures that a deformed protein backbone looks like a protein backbone
Forward model and data fidelity
In the end: Each image with one projection and forward model \(\mathcal{F}_j, j = 1,\ldots, M\)
We use NCC as a similarity measure, i.e., \(\mathcal{L}_y(y') = \sum_i \alpha_i(1-\frac{\langle y_i,y_i'\rangle}{\|y_i\|\|y_i'\|})\)
Protein gradient flow derivation
Group action \(\Phi\) by element-wise multiplication,
momentum map is angular momentum in each component
J(w,p)_i = w_i \times p_i
\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 micrographs,
The adjoint is easy to derive but results in rather unaesthetic expression
Gradient flow
3D reconstruction results
210 residues, ADK, one chain (monomer)
3D reconstruction results
More projections = Good
Less noise = good
3D reconstruction results
Reconstruction to multimer proteins 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
Tomography setup
- 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
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
Runs in a few seconds (but quite small images)
Tomography results
But now turn up the noise... (a factor 100, so not super realistic)
Reconstruction quality starts to deterioriate!
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
Fast and easy to converge
Comparable reconstructions to LDDMM
Easy to regularize
(but this is a story for another time)
Oberwolfach S7
By Erik Jansson
