Diffeomorphic Template Transformers

1 INTRODUCTION . The success of Convolutional Neural Networks ( CNNs ) in many modeling tasks is often attributed to their depth and inductive bias . An important inductive bias in CNNs is spatial symmetry ( e.g . translational equivariance ) which are embedded in the architecture through weight-sharing constraints . Alternatively , spatial transformers constrain networks through predicted spatial affine or thin-platespline transformations . In this work , we investigate a special type of spatial transformer , where the transformations are limited to flexible diffeomorphisms . Diffeomorphisms belong to the group of homeomorphisms that preserve topology by design , and thereby guarantee that relations between structures remain , i.e . connected ( sub- ) regions to stay connected . We propose to use such diffeomorphic spatial transformer in a template transformer setting ( Lee et al. , 2019 ) , where a prior shape is deformed to the output of the model . Here a neural network is used to predict the deformation of the shape , rather than the output itself . By introducing a diffeomorphic mapping of a prior shape , and carefully choosing properties of the prior shape , we can enforce desirable properties on the output , such as a smooth decision boundary or a constraint on the number of connected components . To obtain flexible diffeomorphic transformations , we use a technique known as scaling-and-squaring which has been successfully applied in the context of image registration in prior work ( Dalca et al. , 2018 ) , but has received relatively little attention in other areas in machine learning . In an attempt to increase flexibility of the flow , we try to approximate a time-dependent parameterisation using Baker-Campbell-Hausdorff ( BCH ) formula , rather than a stationary field . Hereby , diffeomorphic constraints are directly built into the architecture itself , not requiring any changes to the loss function . Experimentally , we first validate the diffeomorphic spatial transformer to learn data-invariances in a MNIST handwritten digits classification task , as proposed by ( Jaderberg et al. , 2015 ) to evaluate the original spatial transformer . The results show that better results can be achieved by employing diffeomorphic transformations . Additionally , we explore the use of diffeomorphic mappings in a spatial template transformer set-up for 3D medical breast tissue segmentation . We find that the diffeomorphic spatial transformer is able to deform simple prior shapes , such as a normally distributed energy , into high-quality predictive probability densities . We are successful in limiting the number of connected components in the output and achieve competitive performance measured by quantitative metrics compared to direct estimation of class probabilities . 2 RELATED WORK . Spatial Transformers were introduced by Jaderberg et al . ( 2015 ) as a learnable module that deform an input image , and can be incorporated into CNNs for various tasks . In Spatial Transformer Networks ( STNs ) , the module is used to learn data invariances in order to do better in image classification tasks . The work focuses on simple linear transformations ( e.g . translations , rotations , affine ) but also allows for more flexible mappings such as thin plate spline ( TPS ) transformations . The use of spatial transformations in template transformer setting was first proposed by Lee et al . ( 2019 ) , but does not use diffeomorphisms and requires defining a discrete image as shape prior . In the field of image registration , diffeomorphisms have been actively studied and have been succesfully applied in a variety of methods including LDDMM by Beg et al . ( 2005 ) , Diffeomorphic Demons by Vercauteren et al . ( 2009 ) , and SyN by Avants et al . ( 2008 ) . More recently , efforts have been made to fuse such diffeomorphic image registration approaches with neural networks ( Dalca et al . ( 2018 ) , Haskins et al . ( 2020 ) ) . It is well known that although these models mathematically describe diffeomorphisms , transformations are not always diffeomorphic ; in practice and negative Jacobian determinants can still occur due to approximation errors . To reduce such errors , additional regularisation is often applied ( Bro-Nielsen and Gramkow ( 1996 ) , Ashburner ( 2007 ) , Dalca et al . ( 2018 ) ) , but typically requries careful tuning . Image registration has also been applied to perform segmentation by deforming a basis template commonly referred to as an ’ atlas ’ onto a target image ( Rohlfing et al . ( 2005 ) , Fortunati et al . ( 2013 ) ) , for instance by combining ( e.g . averaging ) manually labelled training annotations ( Gee et al. , 1993 ) . There have been some studies that investigated how to obtain smoother segmentation boundaries in neural-based image registration . For instance , Monteiro et al . ( 2020 ) proposed to model spatial correlation by modeling joint distributions over entire label maps , in contrast to pixel-wise estimates . In other work , post-processing steps have been applied in order to smooth predictions or to enforce topological constraints ( Chlebus et al . ( 2018 ) , Jafari et al . ( 2016 ) ) . There have been some studies that try to enforce more consistent topology during training of neural network , but often use a soft constraint that required alteration of the loss function , such as in Hu et al . ( 2019 ) , and GAN-based approaches which in addition require a separately trained discriminator model Sekuboyina et al . ( 2018 ) . Lastly , there have been some studies in which diffeomorphisms in context of spatial transformer networks were investigated . In Skafte Detlefsen et al . ( 2018 ) , subsequent layers of spatial transformer layers with piece-wise affine transformations ( PCAB ) were used to construct a diffeomorphic neural network , but requires a tessellation strategy ( Freifeld et al . ( 2015 ) , Freifeld et al . ( 2017 ) ) . In Deep Diffeomorphic Normalizing Flows ( Salman et al . ( 2018 ) ) a neural network is used to predict diffeomorphic transformations as normalizing flow but to obtain more expressive posteriors for variational inference . 3 DIFFEOMORPHIC SPATIAL TRANSFORMERS . The Spatial Transformer is a learnable module which explicitly allows for spatial manipulation of data within a neural network . The module takes an input feature map U passed through a learnable function which regresses the transformation parameters θ . A spatial grid G over the output is transformed to an output grid Tθ ( G ) , which is applied to the input U to produce the output O . In the original spatial transformer , θ could represent arbitrary parameterised mappings such as a simple rotation , translation or affine transformation matrices . We propose flexible transformations in the group of diffeomorphisms Tθ ∈ D , which preserve topology , by continuity and continuity of the inverse . In Section 4 , we will describe how we can use a diffeomorphic spatial transformer to warp a shape prior , as illustrated in Figure 1 , in a template transformer setting illustrated in Figure 2 . Diffeomorphic Transformation Let us define the diffeomorphic mapping φ = ψ ( 1 ) v ∈ D using an ordinary differential equation ( ODE ) : ∂ψ ( t ) v ( x ) ∂t = v ( ψ ( t ) v ( x ) ) ( 1 ) where v is a stationary velocity field , ψ ( 0 ) v = Id is the identity transformation and t is time . By integrating over unit time we obtain ψ ( 1 ) v , the time 1 flow of the stationary velocity field v. The most basic way to solve an ordinary differential equation from some initial point x0 is Euler ’ s method , in which the trajectory is approximated by taking small discrete steps and adding the difference to the running approximation in time . The method is straightforward to implement , but may take many steps to converge to good approximations . In this work , we will use a technique known as scaling-and-squaring ( Moler and Van Loan , 2003 ) , which allows for fast exponentiation of stationary velocity fields and thus the solution to the ODE defined in Equation 1 . Scaling-and-Squaring To solve the ODE from Equation 1 , with a stationary velocity field v and the solution is the matrix exponential φ = exp ( v ) , we use is the scaling-and-squaring method ( Moler and Van Loan ( 2003 ) , Arsigny et al . ( 2006 ) ) . The method is very similar to Euler ’ s method , but is typically more efficient by exploiting the relation exp ( v ) = exp ( v/2T ) 2 T with T ∈ N together with the fact that exp ( v ) can be well approximated by a Padé or Taylor approximation near the origin ( i.e . for small ||v|| ) . The main idea is to pick a certain step size T such that ||v||/2T < 0.5 and divide the diagonal values in v by the power integral 2T to obtain the approximation for exp ( v/2T ) ≈ Id+v/2T and then squaring ( self-composing ) it T times to find obtain approximate solution for exp ( v ) . Algorithm 1 : Approximating φ = exp ( v ) using scaling-and-squaring Result : φ = exp ( v ) T ← ceil ( log2 ( max ( ||v|| ) + 1 ) φ0 ← v/2T for t = 1 to T do φt ← φt−1 ◦ φt−1 end The approach can efficiently be implemented in existing numerical differentiation frameworks such as PyTorch or Tensorflow by element-wise dividing the vector components in velocity field v by 2T and then self-composing the resulting field 2T times using the linear grid sampling operation defined in Section 3.1 Spatial Sampling To perform a spatial transformation on the input feature map , a sampler takes a set of sampling points Tθ ( G ) , along with an input feature map U = I with input image I to produce output O . In case of template transformer , explained in Section 4 , the input feature map would be a concatenation U = I ◦ S of an input image with some prior shape S. We follow the general sampling framework described in Jaderberg et al . ( 2015 ) , defined for arbitrary sampling kernels of which the ( sub ) gradients can be computed , and the 3D trilinear interpolation in particular : Oci = H∑ h W∑ w D∑ d Ichwd max ( 0 , 1− |yci − h| ) max ( 0 , 1− |xci − w| ) max ( 0 , 1− |zci − d| ) ( 2 ) This procedure should be differentiable with respect to both the sampling grid coordinates and the input feature map by using partial ( sub ) gradients , allowing it to be used in conjunction with backpropagation . 1For our experiments , we utilized the F.grid_sample function in PyTorch 1.6 to perform grid sampling . Baker–Campbell–Hausdorff formula Instead of parameterising our flow by a single stationary velocity field , we might also think of a piece-wise time-dependent sequence of vector fields . By parameterising the deformation as a time-dependent sequence of velocities we hope improve predictive performance by sequentially modeling larger movements first and detailed refinements thereafter . Composing multiple diffeomorphic transformations will also yield a diffeomorphic transformation , as the space of diffeomorphic transformations D is an algebraic group that is closed under the composition operation . The scaling-and-squaring algorithm offers an efficient way to find diffeomorphic transformations from a stationary vector field , but can not straightforwardly be applied to such time-dependent parameterisations . To address this , we can combine two timepoints , now A and B for simplicity of notation , to form the Lie exponential mapping : exp ( Z ) = exp ( A ) exp ( B ) ( 3 ) and apply the Baker-Campbell-Hausdorff ( BCH ) formula up to a certain order to approximate Z = bch ( A , B ) = ∞∑ n=1 zn ( A , B ) = A+B + 1 2 [ A , B ] + 1 12 [ A , [ A , B ] ] − 1 12 [ B , [ A , B ] ] + · · · ( 4 ) where [ · , · ] is the Lie bracket . We apply the formula to approximate the logarithm of matrix exponentials of two noncommutative velocity fields Z = log ( exp ( A ) exp ( B ) ) and then use scaling-andsquaring one time to find the exponential exp ( Z ) . Binary Tree Composition Naive composition of the T diffeomorphic transformations would result into a long chain of composition operations Φ = ( ( ( ( ( ( φ1◦φ2 ) ◦φ3 ) ◦φ4 ) ◦φ5 ) · · · ) · · ·◦φT−1 ) ◦φT ) . To reduce possible interpolation errors in the resampling from growing as a result of such repetitive composing , we compose the field using a binary tree scheme Φ = ( ( ( φ1 ◦ φ2 ) ◦ ( φ3 ◦ φ4 ) ) ◦ ( · · · ◦ ( φT−1 ◦ φT ) ) ) . Treating the composition scheme as a tree structure , the depth now scales in an order of complexity O ( T ) compared to O ( log ( T ) ) when using naive composition , reducing the maximum number of times an BCH approximation is repetitively applied to a single timepoint .

Summary: The authors present a modification to spatial transformer networks that restricts the transformations to the group of diffeomorphisms. When combined with shape priors, this imposes topological constraints on the mappings produced by the network. These are important considerations in applications such as segmentation tasks where we expect there to be constraints on, for example, the number of connected components. The authors demonstrate the effectiveness of their approach in MNIST experiments and a breast tissue segmentation task.