Bessel Equivariant Networks for Inversion of Transmission Effects in Multi-Mode Optical Fibres
Abstract
We develop a new type of model for solving the task of inverting the transmission effects of multi-mode optical ﬁbres through the construction of an SO + (2 , 1) - equivariant neural network.
This model takes advantage of the of the azimuthal correlations known to exist in ﬁbre speckle patterns and naturally accounts for the difference in spatial arrangement between input and speckle patterns.
In addition, we use a second post-processing network to remove circular artifacts, ﬁll gaps, and sharpen the images, which is required due to the nature of optical ﬁbre transmission.
This two stage approach allows for the inspection of the predicted images produced by the more robust physically motivated equivariant model, which could be useful in a safety-critical application, or by the output of both models, which produces high quality images.
Further, this model can scale to previously unachievable resolutions of imaging with multi-mode optical ﬁbres and is demonstrated on 256 × 256 pixel images.
This is a result of improving the trainable parameter requirement from O ( N 4 ) to O ( m ) , where N is pixel size and m is number of ﬁbre modes.
Finally, this model generalises to new images, outside of the set of training data classes, better than previous models.
Introduction
Multi-mode ﬁbres (MMF) have many potential applications in medical imaging, cryptography and communications.
In the medical domain, the use of multi-mode ﬁbre imaging has potential to create hair-thin endoscopes for imaging sensitive areas of the body.
However, to achieve these applications, the ﬁbre transmission properties must be compensated for to return a clear image (Stasio, 2017).
A
MMF has multiple different ﬁbre modes, each of which propagates at a different velocity.
This leads to an amplitude and phase mixing of the image as it propagates through the ﬁbre (Mitschke, 2016).
As a result, an input image creates a complex-valued speckled pattern on the output of the MMF.
The ability to accurately and in a scalable way learn to invert the transmission effects would unlock
MMF imaging as a useful tool across a range of domains.
This work concerns the use of a single multi-mode ﬁbre and not ﬁbre bundles.
Inverting a speckled image is challenging for multiple reasons.
Firstly, the speckled images have a non-local relationship with respect to the original images.
As a result, solely local patch-based models, such as convolutional neural networks, do not make sense as a solution without some dense mapping function.
Therefore, the non-locality necessitates mapping the speckled images into a spatial arrangement similar to the original images before typical image-based deep learning techniques can be used, such as convolutions and pooling.
In addition, the speckled images have circular correlation, which could be taken advantage of, although as noted by Moran et al.
(2018), ﬁnding these requires solving the inversion, creating a chicken-and-egg problem.
Finally, the ﬁbre is equivalent to an unknown complex transmission matrix (TM) so, the inverse of this could be found using a complex-valued linear model, although this presents challenges in terms of memory requirements.
A mapping between 350 × 350 original and speckled images would result in a TM with 350 4 ≈billion entries, requiring a linear model with as many parameters.
Previous work in inverting the transmission effects of MMFs has either required extensive experimentation to characterise the TM of the ﬁbre ( ˇCižmár & Dholakia, 2011, 2012; Choi et al., 2012;
Mahalati et al., 2013; Papadopoulos et al., 2012; Plöschner et al., 2015; Leite et al., 2021), where the number of experimental measurements required for re-calibration was reduced by Li et al.
(2021) by exploiting sparstiy in the TM; made use of dense linear models (Moran et al., 2018; Fan et al., 2019;
Caramazza et al., 2019); or made use of convolutional models (Borhani et al., 2018; Rahmani et al., 2018).
For the machine learning approaches to tackling the inversion task, those which make use of a dense linear model (Moran et al., 2018; Fan et al., 2019; Caramazza et al., 2019) can naturally account for the difference in spatial arrangement between speckled and original images, although they scale badly with the resolution of the images considered ( O ( N 4 ) for N × N resolution).
On the other hand, in theory, the convolutional neural network (CNN) models (Borhani et al., 2018;
Rahmani et al., 2018) improve upon the scalability issue, but in practice due to the need to approximate the transmission matrix and its non-local effects, the models require a large number of layers in order to be able to effectively map every pixel in the speckled image to every pixel in the original image, and hence in practice do not over come the scalability issues.
All of these approaches are mostly expected to work for classes of objects that belong to the class that was used for training (Borhani et al., 2018), with Rahmani et al.
(2018) making initial steps towards general imaging and
Caramazza et al.
(2019) demonstrating this for a more diverse testing dataset.
We provide further details on each of the previous methods in Appendix A.2.
In this work we present a model which naturally accounts for the difference in spatial arrangement between speckled and original images and scales more efﬁciently than previous methods to higher resolution images.
We believe this is the ﬁrst method to demonstrate an ability to invert 256 ×pixel speckled images into 256 × 256 pixel original images.
Our approach also takes advantage of the circular correlations in the speckled images, and improves upon previous general imaging results.
Concerning the equivariance literature, we develop a model comprising of cylindrical harmonic basis functions, a basis set which has seen little attention in the equivariance literature, and make the connection between the transmission of light through a ﬁbre and the group theoretic understanding used in developing equivariant neural networks.
Our contributions are:
A more data-efﬁcient, scalable model to solve the inversion of MMF transmission effects.
A model that provides better generalisation to out-of-training domain images.
A connection between group theoretic equivariant neural networks and the inversion of MMF transmission effects, providing a new type of model to tackle the problem.
Background 2.1 Multi-Mode Fibres
Multi-mode ﬁbres present a clear advantage over single-mode ﬁbre bundles due to having 1orders of magnitude greater density of modes than a ﬁbre bundle (Choi et al., 2012).
Despite this, the different propagation velocities of each mode, which result in the ﬁbre producing scrambled images, presents a signiﬁcant challenge.
If each propagation mode and velocity was known the
TM could be computed, providing a linear system that inverts the transmission effects, although in general this is not known.
Further, mode speciﬁc losses and imperfect mapping between the input pixels and ﬁbre modes can lead to information loss such that inverting the transmission does not yield the true original input image, as is demonstrated in Figure 1.
Therefore, a model which correctly models the physics of inverting the transmission effects should produce the inverted image in Figure 1 and producing the original image requires addition information to be learned.
We discuss the propagation of light through optical ﬁbres and how we construct theoretical TMs in
Section 2.2 and further details on the inversion of the TM in Appendix A.4.
[Figure]
2.2 Generation of Theoretical Transmission Matrices
Light propagation through a MMF is characterised by the transmission modes of the ﬁbre.
The ﬁbres considered in this work have a step index proﬁle with the refractive index within the core being constant and at a higher value than the cladding.
Analytical solutions for the ﬁbre modes can be determined by solving the Helmholtz equation in cylindrical coordinates (details in Appendix A.1).
The problem can therefore be expressed as an eigenvalue eigenfunction problem, where the eigenfunctions are the ﬁbre modes and the eigenvalues are the propagation constants, β , of the modes.
The solutions for the electric ﬁeld within the ﬁbre are comprised of Bessel functions of the ﬁrst kind inside the core and modiﬁed Bessel functions of the second kind in the cladding as follows:
[Equation]
where r and θ are the radial and azimuthal coordinates, respectively, l is the azimuthal index of the mode, R is the ﬁbre core radius, and u and w are normalised frequencies deﬁned as follows:
[Equation]
where n core and n clad are the refractive indices of the core and cladding, respectively, k 0 is the vacuum wave number, and m is the radial mode index.
Further details on the connection between propagation of light through MMFs and group theoretic equivariant neural networks are provided in Appendix A.1.
Taking the derivative of Equations 1, and equating the resulting functions asserts a smoothness condition on the electric ﬁeld across the core-cladding boundary.
Solving for this for all values of β over all possible integer values of l gives the propagation constants.
These constants correspond to the ﬁbre modes given by the now appropriately parameterised Equations 1.
Examples of the mode ﬁelds are given in Figure 2.
A diagonal ﬁbre propagation matrix, F , can then be made from the eigenvalues, β , of the Helmholtz problem and the corresponding eigenfunctions can be recorded in a matrix, M , which maps the image space to the ﬁbre mode, or group space.
We can ﬁnd the inverse mapping bases from the ﬁbre or group space back to the image space, M † , by taking the conjugate transpose.
These three components allow us to construct the transmission matrix as: TM = MFM † .
2.3 Equivariance
Equivariance in neural networks concerns the study of symmetries and uses these to build an inductive bias into a model (Cohen & Welling, 2016a;
Bronstein et al., 2021).
Equivariance as a property of a neural network guarantees that a transformation of the input data produces a predictable transformation of the predicted features (Worrall et al., 2017).
Formally, we say that, given two transformations T and T ′ which are linear representations of a group G and a network or layer Φ , the network is equivariant if applying transform T g to an input x and then passing it through the network Φ yields the same result as ﬁrst passing x through Φ and then transforming by T ′ g (Cohen & Welling, 2016a); that is,
[Equation]
Convolutional neural networks (CNNs) are an early example of equivariance on images, where by design they are translationally equivariant (LeCun et al., 1998).
More recently, equivariance has been considered for other group actions than that of the translation group.
The majority of works on images involve incorporating rotation or reﬂection equivariance into the model, where a variety of different group transformations have been considered (Cohen & Welling, 2016a,b; Weiler & Cesa, 2019; Murugan et al., 2019;
Wiersma et al., 2020; De Haan et al., 2020; Franzen & Wand, 2021).
We are not interested in developing a group equivariant convolutional model here due to the nonsimilar spatial arrangements between the speckled and original image domains.
Despite this, the concept of equivariance still applies, as the concept of learning a model from a ﬁxed basis set, which guarantees symmetry properties are conserved, is relevant due to the nature of transmission through multi-mode optical ﬁbres.
2.4 Circular Harmonics
Of the works that incorporate rotation equivariance into CNNs we are particularly interested in those using a continuous rotation group SO(2) .
In the case of discrete rotation groups a different feature space is associated to each rotation angle.
Storing features for an inﬁnite number of rotation angles is not computationally tractable.
One option is to turn to a Fourier representation, whereby instead of choosing a discrete number of rotations a maximum frequency can be chosen.
To use a Fourier representation for the weights of the model the domain R 2 is decomposed into a radial proﬁle and angular function.
Solving for the kernel space of permissible ﬁlters that can be used within a CNN and ensuring that the model maintains rotation equivariance yields only the spectrally localised circular harmonics (Worrall et al., 2017; Weiler & Cesa, 2019; Franzen & Wand, 2021).
Therefore a rotation equivariant CNN can be created by solving for the circular harmonics up to a certain rotational frequency, combining this with a radial proﬁle function, and sampling a basis of resolution given by the CNN ﬁlter size.
This yields a set of bases which can be linearly combined with a learnable weighting applied to each to form the kernel for the CNN.
2.4.1 Cylindrical Harmonics
While the circular harmonics have seen some attention in the deep learning community due to the use in constructing rotational equivariant CNNs, the cylindrical harmonics have seen less attention (Klicpera et al., 2020).
The cylindrical harmonics appear as a solution to Bessel functions for integer α and are therefore of interest for problems where information transformation is characterised by such functions.
For example Bessel functions are used when solving wave or heat propagation.
Bessel functions are solutions for different complex numbers α of Bessel’s differential equation:
[Equation]
For integer values of α the Bessel function solutions are a linearly independent set of functions expressed in cylindrical coordinates.
Each function consists of the product of three functions.
The radially dependent term is typically called the cylindrical harmonics.
Further details on the connection between group theory and the Bessel basis functions is provided in Appendix A.1.
Our Method
Our method comprises two models: a Bessel equivariant model and a convolutional post-processing model.
The Bessel equivariant model uses a basis set comprising cylindrical harmonics and learns a complex mapping function to use the known symmetries in optical ﬁbres.
The Bessel equivariant model is constrained by the basis choice to produce circular images, which can have circular artifacts also seen in a true inversion, see Figure 1.
The post-processing model is used to ﬁll gaps in the corners, remove circular artifacts, and sharpen the images.
The two stage approach is useful in safetycritical applications as users can inspect the output of both models, where the Bessel equivariant model output is less likely to depend on the training data due to the choice of inductive bias.
3.1 Bessel Basis Equivariant Model
The core of our model has an inductive bias that better replicates the physics of the task of inverting the transmission effects of MMF imaging.
To achieve this we utilise prior knowledge of MMFs and the modes with which an image propagates through the ﬁbre.
Similar to the successes of rotation equivariant neural networks, which make use of circular harmonics to achieve rotation equivariance, we utilise cylindrical harmonics as an inductive bias of the model.
This is motivated by the propagation modes of the ﬁbre being characterised by Bessel functions, where the radially dependent solution is given by the cylindrical harmonics.
We make the connection to the group theoretic development of equivariant neural networks in Appendix A.1, demonstrating the connection between the group SO
+ (2 , 1) , the light cone, and Bessel function solution to the wave equation.
The ﬁrst stage of the model is the computation of two sets of basis functions.
The ﬁrst set transforms the input image into a function that lives on the group and the second transforms back into the image domain.
The model’s weight matrix linearly transforms the function mapped by the ﬁrst basis set.
The basis ﬁlters are computed in accordance with equations in Appendix 2.2.
These functions are sampled on a grid given by the size of the speckled images for the ﬁrst basis set and the original images for the second basis set.
An example of the bases produced is given in Figure 2.
The Bessel function bases are computed in a pre-training stage and therefore this step only has to be completed once when creating the model.
Further, these bases are not trainable parameters and are not updated during training of the model.
These basis functions are an alternative to those used in general rotation equivariant neural networks, which typically offer equivariance to the group SO(2) , except they correctly model the symmetry group, SO
+ (2 , 1) of the task of inverting transmission effects of a MMF due to the added time dimension.
The model trainable parameters are complex-valued weights of size equal to the number of bases.
The complex weight matrix is a diagonal matrix.
The images transformed by the ﬁrst basis set are multiplied by the weight matrix before being transformed by the second basis set to produce the predicted image.
The overall model architecture is given in Figure 3.
3.2 Post-Processing Model
As demonstrated in Section A.4 the limited number of modes means that the transmission effects of a MMF are typically not fully invertible leading to information loss.
We would therefore not expect an equivariant model informed by the physics of the task to be able to fully invert the transmission effects.
To overcome this we add a second model which is similar to a super-resolution model, except that it does not change the resolution and instead learns to remove circular artifacts, ﬁll in gaps, and sharpen the images.
It is a fully convolutional model comprising of convolutional layers and nonlinearities.
This model takes as input the output of our Bessel equivariant model, which we suspect will be similar to what is achievable by passing the speckled images through an inverted TM, and predicts the original images.
This model is composed of eight convolutional layers with ReLU nonlinearities.
The motivation behind the explicit separation of the inversion and enhancement models is that in some applications, especially safety-critical applications, it is important to be able to see the results of an inversion conditioned on the observations alone, regardless of prior expectations.
Experiments
We investigate performance with both data generated with a theoretical TM and experimental data collected using a real MMF ﬁbre.
We provide details of the theoretical TM generation in Section 2.2.
We use the data from a real ﬁbre, as described in Moran et al.
(2018), to validate that our new approach can model the same data, despite using signiﬁcantly fewer trainable model parameters.
Further, we generate data using a theoretical TM as this allows us to ﬂexibly adapt the parameters of the experiment and create new datasets to validate the new model design.
Finally, the theoretical
TM allows us to experiment with images with higher resolution than has been previously been considered.
We provide details on the training of the models and datasets in Appendix A.3.
4.1 Real Multimode Fibre
We compare our model to the complex-valued linear model developed by Moran et al.
by considering both MNIST and fMNIST data.
For this we train and test separate models for the
MNIST and fMNIST datasets.
In addition, we compare multiple different versions of our model where the diagonal restriction of the mapping between bases is relaxed to allow for a block diagonal structure.
We do this to account for manufacturing defects, sharp bends, dopant diffusion, elliptical cores which causes the diagonal mapping of the ﬁbre matrix within the TM to be block diagonal (Carpenter et al., 2014).
This relaxation allows for x -offsets above and below the main diagonal.
Table 1 shows that our model with a full mapping matrix between bases,
i
[Table]
e. full relaxation of the diagonal constraint, outperforms all other models.
In addition, with a relaxation to allow for a 10 block diagonal structure our model performs comparably.
On the other hand, the complex linear model outperforms our Bessel equivariant model when the diagonal restriction on the mapping function between bases is enforced.
Despite this, our model still provides clear and accurate reconstructions of
[Table]
[Figure]
the target images whilst using orders of magnitude fewer trainable parameters.
Additional results are presented in Tables 6 and 7.
Further to a comparison of the loss values, we visually inspect the reconstruction quality in Figure 4, and we provide further visualisations in Figures 16 and 17.
Thus despite our Bessel equivariant diagonal model achieving a larger loss value, the digit or item is still clearly visible in the prediction and is correctly sharpened to a realistic prediction of the original image by our post-processing model.
The complex linear model can also be seen to predict a realistic looking output of the correct digit and object, although some noise does feature in the prediction.
In addition, we visually compare each of our Bessel equivariant models with the diagonal restriction relaxed in
Figures 10 and 11, which demonstrates that as the diagonal restriction is relaxed the model can better predict the original images with less noise in the prediction from the Bessel equivariant model.
We also compare the number of trainable parameters within the model in Table 2.
This highlights that our model requires two orders of magnitude fewer parameters that the complex linear model when using the diagonal Bessel equivariant model.
Therefore, our model has the potential to scale to higher resolution images where the complex linear model will run out of GPU memory.
As the diagonal mapping between Bessel bases is relaxed to include 10 -block diagonal structure the model still requires two order of magnitude fewer parameters.
In the most ﬂexible version of our model the number of trainable parameters is comparable with the complex Linear model, although we have demonstrated that this level of ﬂexibility is not required to achieve good image reconstruction.
with MNIST or fMNIST data.
[Table]
4.2 Theoretical TM
We compare our model to the complex linear model developed by Moran et al.
(2018) with data created using a theoretical TM.
We create multiple datasets one using 28 × 28 fMNIST images, with 180 × 180 speckled images, another with MNIST images, of the same size, and one with a subset of images from the ImageNet dataset (Deng et al., 2009) where we ﬁx the resolution to be 256 ×of both the images and speckled images.
Using fMNIST and MNIST with a theory TM is useful for developing understanding, due to the ease of creating different experiments and datasets, and for testing the generalisability of the model between the two datasets.
We include the ImageNet-based dataset as these are higher resolution images containing more challenging information than fMNIST.
To the best of our knowledge, demonstrating the ability to invert transmission effects of such high resolution images has not been previously achievable with a machine learning based approach.
[Table]
4.2.1 Generalisability of Models
We provide a comparison of the predictions of a complex valued linear model (Moran et al., 2018), our equivariant model only, and our equivariant model coupled with a post-processing model trained on fMNIST images and tested on both fMNIST and MNIST to analyse to what extent the models can generalise to a new dataset.
Table 3 presents the loss values for both the training and testing of each model.
Considering the fMNIST data this shows that the complex linear and Bessel equivariant models achieve similar loss values.
Although this is a result of the Bessel equivariant model underperforming due to not being able to reconstruct pixel values close to the edge of the image as a consequence of the circular nature of the Bessel bases functions, while the complex linear model under-performs due to failing to reconstruct higher frequency information and over-ﬁtting to the general clothing categories.
Finally, when our Bessel equivariant model is combined with the postprocessing model it outperforms all other models.
The ability to generalise to a new data is assessed through the MNIST dataset.
This shows that our Bessel equivariant model signiﬁcantly out-performs the complex linear model in generalising to a new dataset.
Figure 5 shows reconstructions produced by the three models, with further reconstructions provided in Figure 18.
This demonstrates visually that our model, even without the post-processing part, is able to better reconstruct the original images.
The post-processing model reﬁnes the output of the equivariant model and ﬁlls in gaps outside the ﬁbre.
Despite it being possible to estimate the general category of clothing from the outputs of the complex linear model, the higher frequency information has been lost, i.e. the pattern on the jumper or the curve in the trouser leg.
On the other hand, our model is able to reconstruct this higher frequency information.
Previous works have demonstrated some ability to generalise to new data classes (Rahmani et al., 2018; Caramazza et al., 2019), although in each case the results are not perfect.
Here, we compare the different methods trained on fMNIST and then tested on MNIST data.
Figure 6 demonstrates that our Bessel equivariant model generalises well to a new data domain, with further reconstructions provided in Figure 19.
On the other hand, the complex linear model somewhat predicts the correct digits.
In this out of training domain our post-processing model does not add any value to the original Bessel equivariant model.
This highlights the beneﬁt of our two-stage modelling approach as we have one robust model and a second model to sharpen the images, as a result if a user believes the situation could be unusual they could reliably use the output of our Bessel equivariant model and not the post-processing part.
Further to achieving strong generalisation results, we also consider the effect of noise in the speckled images in Appendix A.6.
Figure 8 shows how the number of trainable model parameters scale.
Even for 28 × 28 images our model requires two orders of magnitude fewer trainable parameters than the complex linear model.
[Figure] [Figure]
.2 Scaling to Larger Images – ImageNet
We now experiment with a subset of images from the ImageNet dataset, where we ﬁx the resolution to be 256 × 256 .
Inverting the transmission effects of such higher resolution images has not been previously achievable.
We only compare our equivariant model with our equivariant model coupled with a fully convolutional post-processing model.
The complex-valued linear model cannot ﬁt in the
GPU memory of an A6000 with 48.7GiB of memory with this resolution of image.
[Figure]
We analyse our model visually, with Figure 7 showing that our Bessel equivariant model produces a reconstruction from which one can determine the type of dog and its activity, although the reconstruction does not capture all the high frequency information, and information in the corners is lost due to the circular nature of the ﬁbre.
Further reconstructions are provided in Figure 20.
When we combine the two models some of the artifacts of the Bessel equivariant model are removed as the post-processing model ﬁlls in information towards the corners and sharpens the image.
We present the loss values for both models in Table 4.
This shows a similar result as for fMNIST, that our Bessel equivariant model can solve the task well, but requires the post-processing model to sharpen the image and ﬁll in gaps in the corners.
We compare the models’ memory requirements in Figure 8.
Note how fewer trainable parameters translate into signiﬁcantly less memory use than the complex linear model.
This effect is seen more drastically when scaling to high resolution images, where the complex linear model runs out of memory on a 24 .
2 GiB Titan RTX for original and speckled images of resolution 180 × 180 pixels.
Our model, in contrast, has been tested on 256 × 256 pixel images, requiring only 2 . 1 GiB, and can scale to megapixel images.
Finally, we explore the impact on each model of reducing the size of the training dataset.
Requiring a reduced dataset size minimises the time taken in a lab collecting data.
Figure 8 demonstrates that our model outperforms the complex linear model with 10 × -less training data, reducing the need to collect 12000 samples to 1200 .
The reconstructions are visualised in Appendix A.12.
[Table]
[Figure]
Conclusions
We present a new type of model for solving the task of inverting the transmission effects of a
MMF
through developing a SO + (2 , 1) -equivariant neural network and combining this with a postprocessing network.
This improves upon complex linear models through incorporating a useful physically informed inductive bias into the equivariant network.
The equivariant network is shown to perform well on new image classes, providing a general model for ﬁbre inversion.
The postprocessing network is speciﬁc to an image domain, as it generalises to new image classes as well as a general convolutional network, although it improves the quality of returned images.
The use of a theoretical TM allows us to compare an ‘ideal’ inverse with the output of previous models based on learning a full transmission matrix.
This suggests that previous learned transmission matrices were combining elements of the actual transmission matrix with elements of the post-processing network, which would affect their ability to generalise to new images.
We also anticipate that in interactive safety-critical applications users might want the ability to switch between modes, to be sure that the evidence was there, and not overly inﬂuenced by the priors in the training data.
This new model signiﬁcantly improves the ability to scale to higher resolution images by improving the scaling law from O ( N 4 ) to O ( m ) , where N is pixel size and m is the number of ﬁbre modes.
We provide a comparison between our model and prior works on both data created using a real ﬁbre and a theoretical TM, demonstrating in both cases that our model solves the task while using signiﬁcantly fewer trainable parameters than the complex and real linear models.
In addition, we demonstrate results on high resolution 256 × 256 images, which has previously been unachievable due to the growth of parameters with previous models.
Furthermore, we demonstrate the ability of our model to generalise to new data classes outside of the training data, outperforming prior works.
The dramatic reduction in the number of parameters for each ﬁbre conﬁguration opens the way for future models which can learn mappings for high-resolution images, from a wider set of perturbed ﬁbre poses and combine these using architectures such as VAEs.
Acknowledgments and Disclosure of Funding
JM is supported by a University of Glasgow Lord Kelvin Adam Smith Studentship.
RM-S,
MP, DF and SPM are grateful for EPSRC support through grants EP/T00097X/1 and RM-S for
EP/R018634/1.
MA is funded by dotPhoton and a UofG scholarship.
DF acknowledges funding through the Royal Academy of Engineering Chair in Emerging Technologies programme.