Moving Frame Net: SE(3)-Equivariant Network for Volumes
Editors: Sophia Sanborn, Christian Shewmake, Simone Azeglio, Arianna Di Bernardo, Nina Miolane
Abstract
Equivariance of neural networks to transformations helps to improve their performance and reduce generalization error in computer vision tasks, as they apply to datasets presenting symmetries (e.g. scalings, rotations, translations).
The method of moving frames is classical for deriving operators invariant to the action of a Lie group in a manifold.
Recently, a rotation and translation equivariant neural network for image data was proposed based on the moving frames approach.
In this paper we significantly improve that approach by reducing the computation of moving frames to only one, at the input stage, instead of repeated computations at each layer.
The equivariance of the resulting architecture is proved theoretically and we build a rotation and translation equivariant neural network to process volumes, i.e. signals on the 3D space.
Our trained model overperforms the benchmarks in the medical volume classification of most of the tested datasets from MedMNIST3D.
1. Introduction
There is currently a great interest in building machine learning methods that respect symmetries, such as translation, rotation and other physical gauge symmetries.
Convolutional neural networks (CNN) are translation equivariant neural networks that have shown great success in a wide variety of tasks related to image processing and understanding.
Recent work has shown that designing group-equivariant CNNs that exploit additional symmetries via group convolutions has even further increased their performance (Cohen and Welling, 2016;
Worrall et al., 2017; Weiler et al., 2018;
Cohen et al., 2019; Bogatskiy et al., 2020).
The method of moving frames (Cartan, 1935; Fels and Olver, 1999), initially proposed by ´Elie Cartan to produce differential invariants, was recently applied to define SE(2)- equivariant convolutional neural networks, i.e. CNNs equivariant to rotations and translations in 2D, by Sangalli et al.
The network constructed with differential invariants provides an alternative to group convolution when constructing group equivariant networks.
In this work, we build on the recent work of Sangalli et al.
(2022) to develop a SE(3)- equivariant neural network.
We propose a new approach that uses differential invariants on 3D CNNs by computing a moving frame from the input of a neural network and applying it to Gaussian n -jets in order to obtain equivariant architectures for SE(3).
This technique is leveraged to propose a novel CNN architecture that we call SE(3)-Moving Frame Network (SE3MovFNet) 1 .
The new architecture can be applied to volumetric data.
Empirically, we show that SE3MovFNet improves the performance of competitive CNNs on a collection of datasets for 3D medical image classification problems.
The paper is organized as follows: We discuss related work in the literature in Section 2 and we introduce concepts of group equivariance, as well as some basics on the method of moving frames, in Section 3.
In Section 4 we derive a moving frame in SE(3) and show how one can obtain an equivariant neural network architecture from a moving frame and in Section 5 we introduce the SE3MovFNet based on the moving frames method
In Section 6 we validate de SE3MovFNet in a task of medical volume classification and overperform most of the benchmarks.
We end the paper with concluding remarks in Section 7.
2. Related Work
In the literature of group-equivariant networks there exist many approaches to plane rotationequivariant networks, for example, Cohen and Welling (2016); Worrall et al.
; Weiler et al.
Also on 2D rotation-equivariant networks, some approaches are based on differential operators (Shen et al., 2020; Jenner and Weiler, 2022; Sangalli et al., 2022).
In particular, the current approach is an extension of the moving frames-based SE(2)- equivariant neural network in Sangalli et al.
In the domain of 3D CNNs, two of the possible data representations are point clouds and volumetric data.
Many approaches that seek equivariance to space rotations are for CNNs that process point cloud data (Thomas et al., 2018;
Chen et al., 2021; Melnyk et al., 2021; Thomas, 2020).
Our work focuses on defining SE(3)-equivariant networks for data based on voxels, i.e., volumetric data.
Some other approaches that aim to achieve this result are: Worrall and
Brostow (2018) achieves equivariance to a discrete subgroup of SO(3); Weiler et al.
(2018) uses a steerable filter basis based on spherical harmonics to learn general SE(3)-equivariant filters and Shen et al.
(2022) does the same thing using filters based on partial differential operators.
Our approach uses differential operators like Shen et al.
but instead of using a steerable filter basis we apply a moving frame to invariantize the network, which consists of evaluating each neighborhood with a rotation computed at the first layer.
3. Technical Background
In this section we introduce both the concepts of group action and equivariance and the basic concepts behind the method of moving frames.
The final goal of the paper is to propose a class of group-equivariant networks based on the method of moving frames.
3.1. Group Actions and Equivariance
Given a group G and a set M , a group action 2 of G on M is a map π : G × M → M such that ∀ x ∈ M π ( e, x ) = x where e is the neutral element of G and ∀ g, h ∈ G , ∀ x ∈ M π ( g, π ( h, x ))
=
π ( g · h, x ).
We denote for all g ∈ G , x ∈ M , g · x :=
π
( g, x ).
If G is a Lie group, M a smooth manifold and π is a smooth map, then π is a Lie group action.
See
Appendix A for properties of group actions.
In this paper we have an implicit assumption of locality of the group action.
A local Lie group action is a smooth map π : U → M where
U ⊆
G × M is an open set such that { e } × M ⊆ U , satisfying ∀ x ∈ M , π ( e, x ) = x and ∀ g, h ∈ G , x ∈ M s.t.
{ ( h, x ) , ( g, π ( h, x )) , ( g · h, x ) } ⊆
U , we have π ( g, π ( h, x ))
=
π ( g · h, x ).
We use the same notation as group actions for local group actions.
Given sets M and N acted upon by G an operator ϕ : M → N is equivariant if ∀ g ∈ G , x ∈ M , ϕ ( g · x ) = g · ϕ ( x ).
We assume that the action on M is not the identity to avoid trivial cases.
Invariance is a special case of equivariance where the action on N is the identity, i.e. ϕ ( g · x ) = ϕ ( x ).
Given manifolds X and Y , when G acts on X with actions we define the action on the space of smooth functions C ∞ ( X, Y ) as, for all f ∈ C ∞ ( X, Y ), x ∈ X , g ∈ G
[Equation]
As this paper is focused on exploring equivariant networks on signals, this is the type of action we seek equivariance to.
Since we are interested in rotations in the input domain we do not consider actions that change the output of the function f , but the general method presented in this paper is capable of dealing with that.
3.2. The Method of Moving Frames
Moving Frames.
Let M be an m -dimensional smooth manifold and G be an r -dimensional
Lie group that acts on M .
A moving frame (Fels and Olver, 1999) is a G -equivariant map ρ : M → G which in particular satisfies, ∀ z ∈ M , g ∈ G [Equation]
[Equation]
A moving frame ρ induces the function z
(cid:55)→ ρ ( z ) · z which is constant over each orbit
O z = { g · z , g ∈ G} .
Namely, ∀ z ∈ M , g ∈ G , ρ ( g · z ) · g · z = ρ ( z ) · z .
Invariantization.
The main interest of having a moving frame from the perspective of equivariant deep learning is the invariantization it defines.
Given an operator F : M → N , its invariantization is defined as ∀ z ∈ M , ı [ F ]( z ) :
= F ( ρ ( z ) · z ).
The invariantization of an operator is invariant with respect to the group action as ı [ F ]( g · z ) =
ı [ F ]( z ) for every z ∈ M , g ∈ G .
Applying the invariantization to an invariant operator returns the same operator, therefore the set of invariant operators is the set of invariantized operators.
In our case, objects of interest (volumes, images, etc) are functions f : X → Y where
X =
R p and Y =
R q , p, q ∈ N ∗ .
They can be modeled as submanifolds of the manifold
M = X
× Y by identifying them by their graph where each point has coordinates ( x , u ) =
( x , f ( x )).
In that case, if we can decompose the action of G into an action on X and an action on Y , then we can associate each invariant operator on M to an equivariant one on the space of functions f : X → Y (see appendix B).
We use this framework in this paper.
Cross-Section.
A cross-section to the group orbits is a submanifold K ⊆ M of dimension complementary to the group dimension i.e. dim K = m − r that intersects each orbit transversally 3 .
If the intersection happens at most once it is a regular cross-section.
If G acts freely and regularly on M and given a regular cross-section K to the group orbits, then for each z ∈ M there is a unique element g z ∈ G such that g z · z ∈ K .
The function ρ : M → G mapping each z to g z is a moving frame (Fels and Olver, 1999; Olver, 2007).
Jet-Bundle.
The n -th order jet bundle (Olver, 1993), or jet-space , J n ( M ) is an extension of a manifold M given by equivalence classes of functions.
For us the jet bundle is particularly useful when the group action is not free on M , as prolonging the manifold to a sufficiently high-order jet bundle and extending the group action to this space can result in a free action, enabling the definition of a moving frame.
In this section we define the jet-space for spaces of the form M = X
× Y where X
= R p , Y = R .
Given a multi-index I
= ( i 1 , · · · , i p ) ∈ N p , we will note | I | =
(cid:80) pk
=
i k its modulus, and let us denote by I pn =
{ I ∈ N p , | I | ≤ n } the set of multi-indices in N p of modulus at most n ∈ N .
For x ∈ X and n ∈ N , we introduce J ( n ) x
f the operator mapping C ∞ ( X, Y ) to Y ( n ) p =
R ( n + pn ), and defined for any f ∈ C ∞ ( X, Y ) by J ( n ) x ( f ) = (cid:0) ∂ I f ( x ) (cid:1) I ∈I pn
=
(cid:0) ∂ i 1 ∂ i 2 . . .
∂
i p f ( x ) (cid:1) ( i 1 ,...,i p ) ∈I pn .
Then for any u ( n ) = (cid:0) u I (cid:1) I ∈I pn ∈ Y ( n ) p and any x ∈ X , ( J ( n ) x ) − 1 ( u ( n ) ) is an equivalence class for the equivalence relation f 1 ∼ f 2 ⇐⇒ J ( n ) x ( f 1 ) =
J ( n ) x ( f 2 ) .
This class is represented in particular by the polynomial function defined for any t ∈ X by
[Equation]
with I !
=
i 1 !
i 2 ! . . .
i p !.
It is the Taylor polynomial of order n at x of any function of the class.
The n th-order jet space of M , noted J n ( M ), is the union of all such equivalence classes, and can therefore be indentified to X × Y ( n ) p .
According to the above, for an element ( x , u ( n ) )
=
( x , ( u I )
I ∈I pn ), each u I is also a partial derivative of u evaluated in x ,
namely u I =
∂
I u ( x ).
For example, if p = 3 and x = ( x, y, z ), u (0 , 0 , 0) =
u ( x ), u (1 , 0 , 0) =
∂u∂x ( x ) =
u x ( x ), u (1 , 1 , 0) =
∂ 2 u ∂x∂y ( x ) =
u xy ( x ) and so on.
In practice we will often use these partial derivative notations to identify elements of the jet space, and omit the variable as it is explicit from the first component.
For example in the case p = 3, n = 2, an element z ∈ M =
J 0 ( M ) is identified by z = ( x , u ) =
( x, y, z, u ) and an element z (2) ∈
J 2 ( M ) by z (2) =
( x , u (2) ) =
( x, y, z, u, u x , u y , u z , u xx , u xy , u yy , u xz , u yz , u zz ).
Prolongation of the Group Action.
Because smooth functions at a point x can be associated to an element of the jet-space and vice-versa, it makes sense that the actions on functions induces an action in the jet-space.
This action is computed by associating an n -jet to a function and computing its derivatives at the transformed point g · x .
Formally, given a point ( x , u ( n ) ) ∈
J n ( M ), let u ∈ C ∞ ( X, Y ) be a function such that u ( n ) =
J ( n ) x u , without loss of generality we can choose u to be the polynomial (3).
We define the prolongation of the action of G on M to the jet-space J n ( M ), given by, for g ∈ G and z ( n ) =
( x , u ( n ) ) ∈
J n ( M ) [Equation]
The expression (4) is well defined as it can be verified that it does not depend on the choice of u Olver (1993) and is a group action. The intuition behind evaluating the derivatives at the point g · x , is that the value in u at x is the same as the value of g · u at g · x , i.e.
( g · u )( g · x ) = u ( g − 1 · g · x ) = u ( x ), however its n -jet is not the same (see Figure 1).
Fundamental invariants.
Invariantizations of operators in the jet-space are referred to as differential invariants and the invariants 4 ı
[ x i ] ı
[ u I ], 1 ≤ 1 ≤ p , 0 ≤ | I | ≤ n are called fundamental invariants because every differential invariant of order n can be expressed as a functional combination F ( ı [ x 1 ] , . . .
, ı
[ x p ] , ( ı [ u I ])
≤| I |≤ n ) (Olver, 2007).
Conversely every function of the fundamental invariants is a differential invariant.
4. Moving Frame and Differential Invariants for SE(3) on Volumes
The group SE(3) of rotations and translations in dimension three, is the semi-direct product of Lie groups SO(3) (rotations) and R 3 (translations) and since both are 3-dimensional, then
SE(3) is a 6-dimensional Lie group.
Here a volume refers to a signal on a 3-dimensional
Euclidean domain, i.e. functions of the type f : R 3 → R q .
We derive SE(3)-equivariant operators on volumes using the method of moving frames.
Volumes are represented as submanifolds of M = R 3 × R .
We consider the case q = 1, but keeping in mind that for higher dimensions it is just a matter of channel-wise application.
SE(3) acts on M by rotating and translating the spatial coordinates x , i.e.
[Equation]
If we proceed to extend M to the first-order jet space we will find that SE(3) does not act freely on J 1 ( M ).
Indeed, the orbit of a point ( x , u, u x , u y , u z ) is the Cartesian product of
R 3 , { u } and a sphere with radius (cid:113) u
2 x + u 2 y + u 2 z , hence it has dimension 5 ̸ = dim SE(3) = 6.
Therefore it is necessary to prolong the action to the second order jet-space in order to be able to obtain a moving frame.
In this section we use a matrix notation for compactness: we denote ∇ u = [ u x , u y , u z ] T and
[Equation]
In that way, the coordinates of the second order jet-space are identified by z (2) =
( x , u (2) ) =
( x , u, ∇ u, Hu ). We compute the prolonged action following (4). Choosing some u ∈ C ∞ ( X, Y ) such that u ( n ) = J ( n ) x u the action becomes ( R, v ) · ( x , u ( n ) ) =
( R x + v , J ( n ) R x + v (( R, v ) · u )).
In order to compute the second order jet-space, we compute the gradient and Hessian matrix of the function ˜ u ( t ) =
(( R, v ) · u )( t ) at the point ˜ x = R x + v . From (1) we have ˜ u ( t ) = u (( R, v ) − 1 · t ) =
u ( R T ( t − v )), thus applying the chain rule we have ∇ ˜ u ( t )
=
R ∇ u
( R T ( t − v )) and substituting t by ˜ x
we have ∇ ˜ u (˜ x )
=
R ∇ u ( x ).
An example illustrating this in the twodimensional case is shown in Figure 1.
A similar argument can be applied to show that [Equation] A similar reasoning can be applied to describe the coordinates of the higher order jet-spaces as symmetric tensors and obtain the prolongation of the group action of higher order using tensor contraction.
This action decomposes into an action on X and an action on Y ( n ) .
Now we can find a cross-section that will give us a moving frame.
Because the crosssection has to have a dimension complementary to the group, we use six equations to construct it.
We also add some inequalities to ensure regularity.
We propose K = { x = 0 , u xy =
u xz =
u yz = 0 , u xx ≥ u yy ≥ u zz , u x > 0 , u y > 0 } i.e. x = {+
+} and Hu is diagonal with its diagonal sorted in non-increasing order.
Hu is symmetric so we can find a P ∈ SO (3)such that P HuP T is diagonal, resulting in the moving frame ρ ( x , u ) ∈ (cid:55)→ ( P, − P x
∈ 4.1. Equivariant Network from a Fixed Moving Frame
From the prolonged group action (7) we can deduce that the non-trivial fundamental invariants of the jet-space of order two are u , v Ti · ∇ u and λ i , i = 1 , 2 , 3, where the v i s are the eigenvectors of Hu (columns of P ) and the λ i s are the eigenvalues of Hu (diagonal coefficients of P HuP T ).
Following the approach of Sangalli et al.
(2022), a two-layer SE(3)-equivariant neural network can be obtained using invariants of order two as follows:
Let the volume f ∈ C ∞ ( X, Y ) =
C ∞ ( R 3 , R ) be the input to the network.
First we compute for all x , J (2) x
f = ( f ( x ) , ∇ f ( x ) , Hf ( x )), followed by the computation of the fundamental invariants of order 2, which we will denote ı
[ z 0 x ] =
ρ ( x , J (2) x f ) · ( x , J (2) x f ) = ρ (
z 0 x ) · z 0 x where z 0 x = ( x , J (2) x f ).
Let ψ 1 : J n ( M ) →
Y 1 , where Y 1 =
R q , be a smooth map.
In a deep learning context we assume ψ 1 to be a multilayer perceptron (MLP).
The first layer ϕ 1 : C ∞ ( R 3 , Y ) → C ∞ ( R 3 , Y 1 ), Y 1 = R q , is given by
[Equation]
and it is SE(3)-equivariant, because ı [ ψ 1 ] is an invariant applied as in Appendix B.
We build the second layer analogously.
The output of the first layer is a signal f 1 = ( f 11 , f 21 , . . .
f q 1 ) =
[Equation] f ) : X → Y 1 .
We compute the derivatives J (2)1 , x f =
( J (2) x ( f 11 ) , . .
, J (2) x
( f q 1 )) and the fundamental invariants ( ı [ z 1 , 1 x ] , · · · , ı [ z 1 ,q x ])
where z 1 ,j x = ( x , J (2) x ( f j 1 )) and ı [ z 1 ,j x ] = ρ ( z 1 ,j x ) ·
z 1 ,j x . Let ψ 2 : ( J n ( M ))
q → Y 2 be a function given by an MLP.
The second layer can be computed from the output of the first by
[Equation]
Again this function is equivariant because it is a function of invariants.
The second layer ϕ 2 = ϕ ′ 2 ◦ ϕ 1 is equivariant because it is a composition of equivariant operators.
This process can be repeated to obtain L equivariant layers ϕ l = ϕ ′ l ◦ ϕ l − 1 , 0 ≤ l ≤ L .
Using the cross-section K , the approach described above requires the computation of the gradient of eigenvectors and eigenvalues with respect to the matrix entries, and even when using a closed polynomial expression to write these values, it can be quite challenging numerically.
With both the numerical or closed form expression of the eigenvectors, the training of the networks resulted in exploding gradients in our early experimentation.
We propose a new solution which limits the computations to only one moving frame.
The alternative we propose is the following.
Instead of computing the differential invariants at each layer, involving the computation of the moving frame based on the previous layer’s feature maps, we compute the moving frame based only on the network input signal and compute all subsequent layers based on this moving frame.
Computing a two-layer network as in the previous example, f 1 is obtained exactly as in step 1.
Now from f 1 we compute J (2)1 , x f for all x .
Given some ψ 2 : ( J n ( M ))
q → R q ′ (which again should be regarded as an MLP)
we can obtain the output of the second layer.
In contrast to (9), however, we transform ρ ( z 1 ,j x ) according to ρ ( z 0 x ) =
ρ ( x , J (2) x f ), not to itself obtaining
[Equation]
The next result shows that repeated application of (10) defines a SE(3)-equivariant network.
Proposition
Let X =
R 3 , Y = R q 0 =
R and Y l = R q l for 1 ≤ l ≤ L , assume that
SE(3) acts on X × Y l like (5) .
Let ρ : X × Y ( n ) → SE (3) be a moving frame.
Let smooth maps ψ l ∈ C ∞ ( J n ( M ) q l − 1 , Y l ) for 1 ≤ l ≤ L .
The functions ϕ l : C ∞ ( X, Y ) → C ∞ ( X, Y l ) , 1 ≤ l ≤ L defined by, for all f ∈ C ∞ ( X, Y ) , x ∈ X , denoting z 0 x = ( x , J ( n ) x f ) and z l x = ( x , J ( n ) x
ϕ l ( f ) j ) 0 ≤ j ≤ q l , 1 ≤ l ≤ L ,
[Equation]
11) and, for 1 < l ≤ L either
[Equation]
or (assuming q l = q l − 1 )
[Equation]
are SE(3)-equivariant for all 1 ≤ l ≤ L .
Where the of G on J n ( M ) q l , 1 ≤ l ≤ L , is the action on J n ( M ) applied coordinate-wise.
Proof See Appendix C.
Because the moving frame is fixed we do not have to compute the gradients of an eigen decomposition, as the moving frame can be seen as an input to the network we only have to compute the gradients of an eigen decomposition once, and the moving frame can be seen as an input to the network.
Indeed the expressions ρ ( z 0 x ) · J ( n ) x ϕ l [ f ] are linear with respect to the values of of the l -th layer ϕ
l [ f ].
5. Moving Frame Nets for SE(3) Acting on Volumes 5.1. Gaussian Derivatives
In order to compute the differential invariants we use Gaussian derivatives.
Gaussian derivatives have already been used in neural networks to produce structured receptive fields (Jacobsen et al., 2016; Penaud-Polge et al., 2022; Sangalli et al., 2022) in CNNs.
Gaussian derivatives are used to compute the derivatives of a Gaussian filtered volume f : Ω → R defined on a grid Ω ⊆ R 3 :
[Equation]
We refer to the Gaussian n -jet of a volume as the Gaussian derivatives of order ≤ n f ( n ) σ :=
( f ∗ G i,j,kσ ) 0 ≤ i + j + k ≤ n .
We can also identify the Gaussian n -jet by tensor coordinates.
In particular for n = 2 we write f ( n ) σ = ( f σ , ∇ σ f, H σ f ) where f σ = f ∗ G σ , ∇ σ is the Gaussian gradient i.e. Gaussian derivatives of order one and H σ f is the Gaussian Hessian, i.e. Gaussian derivatives of order two.
Given a orthogonal matrix for each point P : Ω → SO (3) (e.g. the matrices defining a moving frame) we denote the local prolonged action by ( P · f ( n ) σ )( x ), e.g. for n = 2 ( P · f ( n ) σ )( x ) =
( f σ ( x ) , P ( x ) ∇ σ f ( x ) , P ( x ) H σ f ( x ) P ( x ) T ).
If f : Ω → R q is a multi-channel volume we can apply these operations channel-wise.
Gaussian filters are already rotation-equivariant, so their composition with a differential invariant yields a rotation-equivariant operator.
Moreover, they avoid the issues inherent with discrete signals and reduce the negative impact sampling signals.
These properties motivate the use of Gaussian n -jets to compute the invariants.
5.2. Architecture
Based on the exposition on Section 4.1, the general idea of our SE(3)-equivariant architecture, given an input signal f : Ω → R where Ω ⊆ R 3 is a three-dimensional grid, is to first compute the matrices P : Ω → SO (3) of the moving frame diagonalizing H σ ′ f
( x ) for every x , i.e., find P ( x ) such that P ( x ) H σ ′ f
( x ) P ( x ) T .
Even if all eigenvalues are different there are at least two choices of normalized eigenvectors (i.e. columns of P ( x )) corresponding to each eigenvalue, therefore to remove ambiguity an keep a consistent moving frame, we choose the option that has smallest angle with the gradient, and if the gradient norm is too small we multiply that column by zero.
After computing P we compute blocks as shown in Figure 2( b ), which we call SE3MovF blocks, using the moving frame and the current features maps as input.
The scale of each layer does not need to be necessarily the same, here we consider that a scale σ ′ is used to compute the moving frames and a scale σ to compute the derivatives at each block.
A simple form of a global architecture, which we call SE3MovFNet, is in Figure 2( a ).
The feature maps of each block are summed like in residual networks, which mimics a PDE scheme (Ruthotto and Haber, 2020).
The network in Figure 2 is specialized for a fixed number of channels, but by applying an 1 × 1 × 1 convolution between blocks we can increase the number of feature maps of the next layer.
Pooling may also be performed by subsampling after a block.
The global max-pooling at the end renders the equivariant architecture invariant (Bronstein et al., 2021), which is interesting for a classification problem.
[Figure]
P
S E
M o v
F B l o c k
S E
M o v
F B l o c k +
S E
M o v
F B l o c k +
S E
M o v
F B l o c k + · · ·
S E
M o v
F B l o c k +
G l o b a l M a x P oo li n g u out
[Figure]
fundamental build block the SE(3) MovF block.
The block P denotes the computation of the moving frame matrices P : Ω → SO (3).
6. Experiments
MedMNIST.
MedMNIST
(Yang et al., 2021a,b) is a collection of datasets for benchmarking algorithms in medical image processings classification tasks.
It contains six datasets of 28 × 28 × 28 volumes: AdrenalMNIST3D, NoduleMNIST3D, VesselMNIST3D, SynapseMNIST3D, OrganMNIST3D, FractureMNIST3D.
For more information see Yang et al.
(2021b).
For each dataset we train a network with five SE3MovFr blocks with 16 , 16 , 32 , 32 ,filters, using a stride of two in the second block.
Voxelwise MLPs are computed as two subsequent 1 × 1 × 1 convolutions followed by batch normalization (both) and leaky ReLU (only the first) and with the same number of neurons.
We also train a CNN baseline with the same number of filters where each block consists of two 3 × 3 × 3 convolutions followed by batch normalization and leaky ReLU.
Input volumes are resized to 29 × 29 × 29 so that
[Table]
subsampling by a factor of two is equivariant by rotations of 90 ◦ around the coordinateaxes.
Overall results can be seen in Table 1.
There we can see that the SE3MovFNet surpassed most of the benchmarks.
Results of testing the models on rotated test sets are seen in Figure 3 and
Appendix E.
In those results we observe it has perfect invariance for 90 ◦ rotations, evidenced by the periodicity of results, and a generally better equivariance than the CNN baseline with or without augmentation.
It suffers, however, a significant loss for orientations not multiple of 90 ◦ .
[Figure]
7. Conclusions
We have developed and successfully applied a SE(3)-equivariant architecture, SE3MovFNet, for a classification task in medical image processing.
The proposed SE3MovFNet is an extension of a previous approach for SE(2)-equivariant networks (Sangalli et al., 2022) that corrects some of its numerical issues.
The performance of our network is overall positive, as it attained the best results in 4 of the 6 evaluated datasets of MedMNIST and maintains a reasonable accuracy when images are rotated.
Future work will explore other symmetries on other manifolds like for example scale and rotation symmetry simultaneously for images or volumes.
Appendix A. Properties of Group Actions
If G is a Lie group acting on a manifold M we define the orbit passing through a point z ∈ M as O z as the submanifold
[Equation]
We can classify the action of a Lie group as
• semi-regular if all the group orbits have the same dimension; • regular if it is semi-regular and each point z ∈ M as an arbitrarily small neighborhood
U containing z such that the intersection of U and each orbit is connected; • free if for g ∈ G , z ∈ M , g · z = z = ⇒ g
= e where e is the identity on G .
• locally free if there exists a neighborhood U
⊆ G containing e satisfying ∀ g ∈ U ,
[Equation]
Appendix B. From Invariants to Equivariant Operators
We assume we are in the context where we compute the invariantization of operators F :
M 1 → M 2 , with M 1 =
X × Y and M 2 =
Z , and the action G can be decomposed into actions in X and Y , i.e.
[Equation]
for all g ∈ G , ( x , u ) ∈ X × Y .
An equivariant operator in M 1 can be related to an equivariant one in the space of functions Y X .
Suppose ψ : M 1 =
X × Y → Z is G -invariant, then take ¯ ψ : Y X
→ Z X to be ¯ ψ [ f ]( x ) = ψ ( x , f ( x )), for all x ∈ X , f ∈ Y X .
Assuming that the action on Y is the identity g · u = u for u ∈ Y and the same for Z we have
[Equation]
therefore g · ¯ ψ ( f ) =
¯ ψ ( g · f ) for all g ∈ G , f ∈ Y X .
In other words, an invariant operator in the Cartesian product X × Y to Z induces an equivariant operator taking functions in Y X to functions in Z X .
The same reasoning can be applied if ψ : J n ( M 1 ) =
X
× Y ( n ) is an invariant on the
Jet-space and ¯ ψ : C ∞ ( X, Y ) → C ∞ ( X, Z ) is the operator ¯ ψ
[ f ]( x ) = ψ ( x , J ( n ) x f ) to show that ¯ ψ is G -equivariant.
We have
[Equation]
Appendix C. Proof of Proposition
Proof For l = 1 we have
[Equation]
Noting y = g − 1 · x and recalling (7), g · ( x , J ( n ) x f ) = ( g · x , J ( n ) ( g · x ) ( g · f )), we can simplify [Equation]
[Equation] Hence, [Equation] by invariance of z
[Equation]
by definition of the action on functions .
[Equation] Therefore, ϕ 1 ( g · f ) = g · ϕ 1 ( f ). Now, for l > 1, we have for the case (12),
[Equation]
As shown earlier, ( x , J ( n ) x ( g · f )) = g · z 0 y with y = g − 1 · x . Similarly, assuming that ϕ l − 1 is equivariant,
[Equation]
[Equation] so that [Equation] Since furthermore ρ ( g · z 0 y ) = ρ ( z 0 y ) · g − 1 by definition of a moving frame, we finally get
[Equation]
As for the case (13),
[Equation]
like above, with y = g − 1 · x [Equation] Therefore in all cases ϕ l ( g · f ) = g · ϕ l ( f ) provided this is true for ϕ l − 1 , and the proposition follows by induction.
Appendix D. Complexity Analysis
Let us assume that the input is given as a signal f : Ω → R where Ω = { 0 , . .
, W − 1 } × { 0 , . .
, H − 1 } × { 0 , . . .
, D − 1 } is a grid of size W × H × D .
Moreover, let us assume that we compute the moving frame using Gaussian derivatives of scale σ ′ and the derivatives at other layers using σ , and that the discrete Gaussian derivative filters all have dimension w ′ × w ′ × w ′ , for the moving frame and w × w × w for the other layers.
The computation of the moving frame is done as follows:
• compute all Gaussian derivatives of order one and two of f .
Gaussian derivatives are separable, thus each one can be obtained by three convolutions with a filter of size w , which have cost a cost of O ( W HDw ′ ) floating point operations (flops); • compute the eigenvectors of the Hessian.
Since the matrices have constant size 3 ×we consider this operation is done in constant time for each pixel and this step is done in O ( W HD ) flops.
So the computation of the moving frame is done in O ( W HDw ′ ) flops.
From there on if we compute a layer with q ′ input feature maps and q output feature maps:
• this layer computes (cid:0) n +3 n (cid:1)
Gaussian derivatives for each input feature map, where n is the order of differentiation used, resulting in O (cid:0)(cid:0) n +3 n (cid:1) q ′ wW HD (cid:1) flops; • to compute the prolonged group action, it can be verified that the equivariant group action can be expressed as polynomial in the partial derivatives, and thus it takes O (cid:0)(cid:0) n +3 n (cid:1) q ′ W HD (cid:1) flops; • the previous step is followed by an L -layer multi-layer perceptron at each voxel.
Assuming that the output dimension at each layer of the MLP is at most q we have that this step takes O ( q ′ q
+ Lq 2 ) flops.
The complexity of a layer of SE3MovF is the sum of the complexity of each step, i.e. it can be done in O (cid:0)(cid:0) n +3 n (cid:1) q ′ wW HD + q ′ q + Lq 2 (cid:1) flops.
In our experiments here we used n =and L = 2 for all models, so the impact of those terms is very limited.
Appendix E. Additional Results
Figures 4 and 5 show some more results on MedMNIST3D.
[Figure]
CNN augmented
SE3MovFNet SE3MovFNet aug. [Figure] ( d ) AdrenalMNIST3D
[Figure]
( g ) SynapseMNIST3D
[Figure]
[Figure]