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)