# Self-attention-based non-linear basis transformations for compact latent
space modelling of dynamic optical fibre transmission matrices

Yijie Zheng1 Robert J. Kilpatrick2 David B. Phillips2 George S.D. Gordon1∗
1 Optics and Photonics Research Group, University of Nottingham, NG7 2RD
2 Department of Physics and Astronomy, University of Exeter, Exeter. EX4 4QL
∗<EMAIL_ADDRESS>
###### Abstract

Multimode optical fibres are hair-thin strands of glass that efficiently
transport light. They promise next-generation ultra-fine medical endoscopes
that provide unprecedented sub-cellular image resolution deep inside the body.
However, confining light to such thin fibres means that images are inherently
scrambled in transit.

Conventionally, this scrambling has been compensated by pre-calibrating how a
specific fibre scrambles light and solving a stationary linear matrix equation
that represents a physical model of the fibre. However, as the technology
develops towards real-world deployment, the unscrambling process must account
for dynamic changes in the matrix representing the fibre’s effect on light,
due to factors such as movement and temperature shifts, and non-linearities
resulting from the inaccessibility of the fibre tip when inside the body. Such
complex, dynamic and nonlinear behaviour is well-suited to approximation by
neural networks, but most leading image reconstruction networks rely on
convolutional layers, which assume strong correlations between adjacent pixels
– a strong inductive bias that is inappropriate for fibre matrices which may
be expressed in a wide range of arbitrary coordinate representations with
long-range correlations. We introduce a new concept that uses self-attention
layers to dynamically transform the coordinate representations of dynamically
varying fibre matrices to a basis that admits compact, low-dimensional latent-
space representations suitable for further processing.

We demonstrate the effectiveness of this approach on diverse fibre matrix
datasets, including both randomly generated and physically modeled perturbed
fibre matrices. We show our models significantly improve the sparsity of fibre
bases in their transformed bases with a participation ratio, $p$, as a measure
of sparsity, of between 0.01 and 0.11. Further, we show that these transformed
representations admit reconstruction of the original matrices with $\leq 10\%$
reconstruction error, demonstrating the invertibility of the transformation.

##  1 Introduction

Lensless endoscopic imaging using ultra-thin multi-mode fibre (MMF) represents
a state-of-the-art approach to achieve minimally invasive in vivo imaging. It
enables access to inner regions of the body, such as the brain or blood
vessels, for advanced biomedical imaging devices [1]. These MMFs efficiently
transmit light along their length due to a tailored refractive index profile,
which is modelled physically by linear differential equations (e.g., Maxwell’s
equations). When spatially discretized, these equations become linear matrix
equations, with input and output fields represented as complex-valued vectors
and the effect of the fibres represented by complex-valued transmission
matrices (TMs):

$\mathbf{E_{out}}=\mathbf{T}\mathbf{E_{in}}$ (1)

where $\mathbf{E_{in}}\in\mathbb{C}^{N}$ is a a discretized and column-wise
stacked vector representing the input field at one end of the fibre,
$\mathbf{E_{out}}\in\mathbb{C}^{N}$ is the equivalent discretized output
field, and $\mathbf{T}\in\mathbb{C}^{N\times N}$ is the complex valued TM of
the fibre. Scalar $N$ represents the number of spatial modes supported by the
fibre, equivalent to the number of independent pixels in transmitted images.

Conventionally, such linear matrix equations have been sufficient to model
fibre TMs, but dynamic effects such as bending, temperature, manufacturing
variation and wavelength modulation that are encountered in realistic usage
introduce dynamic behaviour that requires more complex physical modeling
[2][3]. Trainable neural networks represent an alternative modelling approach
that allows approximation of complex dynamic behaviour, but previous attempts
to model TMs has relied on convolutional neural networks (CNN), originally
designed for image data. CNNs have a strong inductive bias tailored to the
statistical properties of natural images and therefore work to date has
typically only modelled static TMs representing single image-to-image
transformations. However, because TMs are necessarily different to images, and
are often expressed in a range of arbitrary coordinate systems, the inductive
bias of CNNs is not suitable for dynamic TM modeling. We therefore present a
model that exploits self-attention-based neural networks to learn arbitrary
input dependent (i.e., nonlinear) basis transformations of TMs that enable
compact latent space representations, and thus efficient dynamic modeling and
reconstruction.

###  1.1 Optical fibre imaging background

As defined in Eqn 1 a TM, $\mathbf{T}$, is a linear mapping from an input
field to an output field. The TM will change if a fibre is perturbed - and
indeed is very sensitive to the bend configuration or temperature of the
fibre. Furthermore, the statistical properties of the relationship between
matrix elements depends upon the basis representation of the TM (e.g., it will
depend upon the choice of mapping describing how pixels from rectangular
images are stacked into column vectors, or choice of basis representation of
the transmitted optical field $\mathbf{E_{out}}$ – such as canonical or
Fourier component representation).

Generally, a TM can be transformed into a new coordinate space that is better
suited for the particular application. For example, if the eigenvectors of the
fibre’s transmission matrix are known and used as a basis, the new TM becomes
purely diagonal [4]. This, in turn, enables significantly more memory-
efficient storage, requiring storage of only the $N$ diagonal elements (in
contrast to $N^{2}$ elements for a full matrix) for an $N$-pixel image
resolution. In addition, an understanding of a sparse fibre TM representation
substantially reduces the number of TM calibration measurements required, as
all diagonal TM values can be obtained with only a single measurement [3] (in
contrast to $N$ calibration measurements needed for a full matrix).

Conventional fibre calibration methods involve learning static TMs by
measuring input-output pairs and solving linear matrix equations to find the
TM. Numerous variants have been demonstrated that achieve image reconstruction
by either explicitly or implicitly utilising the static complex TM including
wavefront modulation [5, 6, 7, 8, 9, 10, 11], speckle-correlated intensity
reconstruction[12, 13, 14, 15] and compressive sensing[16, 17, 18]. Other
advanced computational methods such as Bayesian inference-based algorithm
[19], convex optimization [20], extended Kalman filter [21] and Gradient-
Descent based optimization [22] further reduce computational resources and
time while preserving both amplitude and phase information of the images.

These methods are considered static and linear since they infer a TM based on
input-output pairs of a static fibre. However, in realistic deployments such
as medical imaging, the fibre will be constantly moving and experiencing
temperature fluctuations, thus changing the TM dynamically. Furthermore, the
input-output image pairs needed for calibration will not be obtainable as this
would require placing unacceptably bulky optical components at the tip of the
fibre located deep in the body. Therefore, realistic real-time fibre
characterisation methods have to be able to model a dynamically changing TM
and this model must be updated using only measurements from a single end of
the fibre. Such measurements are typically nonlinear as they involve forward
TM estimation ($\mathbf{T}$) based on round-trip measurements
($\mathbf{T}^{\top}\mathbf{T}$). Previously demonstrated methods include using
feedback from guide stars [23, 24, 8], tracked beacon sources [25, 26], and
reflective structures [27, 28, 29] that can be put onto the tip of the fibre.
To complement this, there are also several physical models designed to model
dynamic TMs under physical perturbations and wavelength modulation [30, 2].
However, these models become intractable when modelling multiple types of
perturbations that must be estimated using as few measurements as possible.
There is therefore a need to develop a more flexible model for dynamic TMs
that could be fitted to measured data.

###  1.2 Neural Networks background

Convolutional Neural Networks (CNN) have been widely used to enable image
transmission and wavefront shaping through mostly static optical fibres [31,
32, 33, 34, 35, 36]. These methods directly reconstruct images and so are
well-suited to CNNs which contain an image-specific strong inductive bias.
Some modifications have been made to better match the physical nature of MMFs
for example the ability to encode the phase information into the networks
where [37] applied complex-weighted, deep neural networks to better represent
speckled images and [38] added a self-attention mechanism to improve the image
quality.

However, most previous methods focus on the calibration of static TMs and aim
to directly reconstruct images. By contrast, dynamic TM calibration that model
a varying TM using neural networks show great promise as they can learn prior
distributions over the perturbed TMs [39] leading to improved generalization.
However, the dynamic changing of the TM (which can be considered a basis
change) can require intractably large numbers of parameters.

Therefore, there are several particular requirements for neural networks used
to model dynamic TMs:

  1. 1.

The model needs to cope with arbitrary representation bases, creating some
degree of basis invariance.

  2. 2.

The model should be sufficiently memory efficient to cope with typical image
resolutions (noting that a fully connected TM requires $N^{2}$ elements for
the $N$-pixel image transmission).

  3. 3.

The model should support complex numbers.

In this paper, to efficiently learn relevant features of ensembles of MMF TMs
under dynamic variation, we apply the ‘Modeling is Compression’ philosophy
[40]. We expect to find a TM representation that is maximally sparse, measured
by a metric such as participation ratio or by the bottleneck size of an
autoencoder, which also meets our criteria of being basis invariant.

Figure 1: Illustration of the various basis transformation methods that can be
performed on a set of perturbed fibre TMs. (a) individual perfect
diagonalization (overfitting). (b) static transformation (underfitting). (c)
input-dependent non-linear regularised transformation (approaching optimal
fitting). The complex-valued colormap shown here is used throughout the rest
of the paper.

Evidentally perfect TM diagonalization as shown in Figure 1(a) is possible by
constructing a unique transformation for each TM, but this is equivalent to a
highly overfitted approach that does not exploit the relative constrained
prior distribution over TM perturbations, or learn any information about how
to generalise to unseen TMs. At the other end of the spectrum, a static
transformation shown in (b) across all perturbations can be learned but will
produce low sparsity since it is underfitted. Our proposed approach shown in
(c) lies between these two extremes and provides an input-dependent but
regularised basis transformation computed using a self-attention architecture,
which can be considered to be approaching an optimal fit.

Generally, self-attention is essential for capturing dependencies between
different elements within a sequence. It has been widely used in natural
language processing such as machine translation [41, 42], text summarization
[43, 44], and sentiment analysis[45], where understanding the context and
relationships between words is crucial for accurate predictions. Beyond
language/word embedding problems, self-attention mechanisms have begun to be
applied to physical systems, for example as generative diffusion models
enabling image reconstructions from on-sensor measurements on a nanophotonic
array camera [46]. In this context, self attention can be considered as a
basis transformation process by allowing each element to attend to all other
elements in the sequence. In our problem, each TM can be considered as a
sentence and each column entry as a word, and we expect self-attention to
enable the model to weigh the importance of each element based on its
relevance to the others, resulting in effective representation learning.

We demonstrate our methods by application to diverse datasets of TMs (i.e.,
randomly-generated forward TMs, randomly-generated round-trip TMs, and
physically modeled TMs under perturbations). We firstly show CNNs have limited
ability to perform the basis transformation of TMs, especially for dense
matrices, and do thus not meet our first criterion. We then show that Fully
Connected Neural Networks (FCNN) are better at transforming TMs to a suitable
basis but scale exponentially with increasing matrix size. Finally, adding
self-attention mechanisms significantly improves on both CNNs and FCNNs in
terms of sparse representation of matrices for all datasets used and offers
better scalability.

We believe this form of TM modelling and data compression could pave the way
for fibre imaging with only optical access to a single-end, enabling the hair-
thin imaging devices that could fit in a needle [29, 39, 3]. More generally,
our approach may find applications to the efficient modelling of matrix data
with other types of scattering media [47, 48]. These methods may also
facilitate real-time tracking of fibre deformation or dynamic scattering media
for deep-tissue imaging, which allows for future expansion of modelling fibre
behavior across wavelengths, e.g., for fluorescence imaging [2].

##  2 Methods

###  2.1 Diverse TM datasets

We generate three diverse datasets (with different sparsity properties) of
$78\times 78$ complex-valued perturbed TMs: randomly-generated forward TMs,
randomly-generated round-trip TMs and physically modeled TMs. For the
randomly-generated forward TMs, $\mathbf{T_{f}}\in\mathbb{C}^{78\times 78}$,
we construct the simulation model [39] to recreate typical properties (e.g.,
high sparsity, high power intensity along the main diagonal and sub-diagonals,
condition number 3-10) observed in real fibre TMs. These are inherently sparse
in their original basis. For the randomly generated round-trip TMs,
$\mathbf{T_{r}}\in\mathbb{C}^{78\times 78}$, based on a nonlinear reflection-
mode fibre imaging system [29, 39], we use the forward TM, $\mathbf{T_{f}}$,
to compute a round-trip TM as:

$\mathbf{T_{r}}=\mathbf{T_{f}}^{T}\mathbf{R}\mathbf{T_{f}}$ (2)

where $\mathbf{R}\in\mathbb{C}^{78\times 78}$ is a matrix that numerically
represents reflection off the distal fibre facet, which contains a highly
scattering reflector. As a result, these round-trip TMs are inherently dense
in their original basis.

Finally, we applied a physical fibre model where each perturbed fibre can be
represented as a series of connected segments, with specifically defined
lengths and bend radii [4]. The TM of the ensemble of segments, when connected
end-to-end, $\mathbf{T_{e}}\in\mathbb{C}^{78\times 78}$, is calculated from
the product of the TMs of the individual segments,
$\mathbf{T_{s_{1}..z}}\in\mathbb{C}^{78\times 78}$:

$\mathbf{T_{e}}=\mathbf{T_{s_{1}}}\mathbf{T_{s_{2}}}...\mathbf{T_{s_{z}}}$ (3)

Here, this dataset comprises TMs with different fibre lengths and bend radius
with fixed core radius of $8\mu$m, numerical aperture (NA) number of 0.22 and
Poisson ratio of 0.17.

For each dataset, we generate 22,000 TMs and split them into training,
validation and test sets with an 8:2:1 ratio. The validation set is expected
to provide unbiased evaluations and stopping criteria and the test set aims to
examine the generalization performance of the model on unseen data.

###  2.2 Basis transformation model

Figure 2: Basis transformation models. (a) Linear similarity transformation-
based model, with (b) a lossless invertible process defined. (c) Non-linear
transformation using different neural network architectures, including CNN,
FCNN, self-attention only and self-attention-based FCNN. Two types of
autoencoder-based constraints are defined for the purpose of (d) sparsity
representation and (e) model invertibility.

Our objective is to compare several different approaches to basis
transformation and thus present a versatile framework for modelling the fibre
TMs in efficient coordinate systems as judged by their sparsity of
representation. We construct and train the models using Tensorflow 2.1 running
on an NVIDIA GeForce RTX 3090 GPU. The Adam optimizer was used with a learning
rate of 0.001 in a decay rate of 1e-5.

####  2.2.1 Linear similarity transformation-based model

We first construct a linear similarity transformation-based model as shown in
Figure2(a), to linearly convert TMs into a new coordinate system based on the
theory of similarity transformation [49]. Specifically, we developed a
differentiable custom layer, where $\textbf{B}\in\mathbb{C}^{78\times 78}$ is
a non-singular square basis matrix that transfers sets of TMs (i.e.
$\mathbf{T_{1...m}}\in\mathbb{C}^{78\times 78}$) to a new basis (i.e.
$\mathbf{T^{\prime}_{1...m}}\in\mathbb{C}^{78\times 78}$), while preserving
the matrices properties (determinant, trace and eigenvalues):

$\mathbf{T^{\prime}_{1...m}}=\mathbf{B}^{-1}\mathbf{T_{1...m}}\mathbf{B}$ (4)

This is a differentiable operation where the model is expected to learn the
basis matrix B during the training, in order to enable a low-dimensional
representation (i.e., increased sparsity) in the new
basis,$\mathbf{T^{\prime}_{1...m}}\in\mathbb{C}^{78\times 78}$, by minimizing
an $\ell_{1}$ norm loss function ($\mathcal{L}_{\mathrm{linear}}$):

$\mathcal{L}_{\mathrm{linear}}=\sum_{i=1}^{m}\|\mathbf{T^{\prime}_{i}}\|_{1}$
(5)

This linear similarity-based transformation model, in turn, ensures a lossless
invertible transformation of TMs to the original basis based on the basis
matrix $\mathbf{B}$ learned by the model during the training.

####  2.2.2 Non-linear transformation models

We next construct non-linear basis transformation frameworks performed with
four different non-linear neural networks models, namely convolutional neural
networks (CNN), fully-connected neural networks (FCNN), single self-attention
layer and self-attention-based FCNN as shown in Figure2. To feed the model, we
convert the input sets of TMs, $\mathbf{T_{1...m}}\in\mathbb{C}^{78\times 78}$
into sets of real-valued TMs $\mathbf{T_{1...m}}\in\mathbb{R}^{156\times 156}$
by separating the real and imaginary parts[50]. We expect to obtain
corresponding sets of TMs,
$\mathbf{T^{\prime}_{1...m}}\in\mathbb{R}^{156\times 156}$, as the output that
are low-dimensional representations (i.e., highly sparse).

The first model is a CNN model that contains four Conv2D and MaxPooling layers
with 190,965 trainable parameters and uses LeakyRelu activations with a
negative slope value defined as 0.3.

The second model is an FCNN comprising two 1D dense layers with the LeakyRelu
activation function (negative slope value defined as 0.3), with a comparable
number of trainable parameters (203,534) to that of the CNN model.

The third and fourth models make use of self-attention mechanisms. Here, we
encode the input information of each TM to two sets of interpretable matrices,
Query ($\mathbf{Q}_{1...j}\in\mathbb{R}^{156\times 156}$) and Key
($\mathbf{K}_{1...j}\in\mathbb{R}^{156\times 156}$), with local interactions
defined within each column. The output of the attention mechanism can be
therefore updated values ($\mathbf{T^{\prime}}_{1...j}\in\mathbb{R}^{156\times
156}$) for each element within the column groups accordingly:

$\mathbf{T^{\prime}}_{1...j}=\mathrm{softmax}\left[\frac{\mathbf{Q}_{1...j}\mathbf{K}^{T}_{1...j}}{\sqrt{d_{k}}}\right]\mathbf{T}_{1...j}$
(6)

The third model directly learns this self-attention transformation, with the
softmax as the only nonlinearity. The fourth and final model adds an FCNN
layer following the self-attention block to enhance the modeling capability.

####  2.2.3 Autoencoder-based constraint

By optimising these non-linear transformation models using only the $\ell_{1}$
norm, we may inadvertently lose inherent properties of the matrices. For
example, the model might produce a transformation that maps all input values
to zero.

To address this, we introduce an invertibility constraint, as shown in Figure
2(e), using an Autoencoder structure to reconstruct the original matrices
(i.e., $\mathbf{T}_{1...m}\in\mathbb{R}^{156\times 156}$) from the transformed
TMs (i.e. $\mathbf{T^{\prime}}_{1...m}\in\mathbb{R}^{156\times 156}$). This
constraint aims to preserve the matrix structure during training by minimizing
the reconstruction loss, $\mathcal{L}_{2}$.

We also enforce an additional sparsity constraint using a second Autoencoder
model as shown in Figure 2(d), where the Encoder is designed to maximally
compress the transformed matrices
($\mathbf{T^{\prime}_{1...m}}\in\mathbb{R}^{156\times 156}$) into a low-
dimensional latent space (LS1) via a ‘bottle neck’ layer. Training this second
Decoder ensures the LS1 contains the necessary information to reconstruct the
transformed TM and is optimized by minimizing the reconstruction loss,
$\mathcal{L}_{3}$. Therefore, the total loss of the whole framework can be
defined as:

$\mathcal{L}_{\mathrm{total}}=\underbrace{\alpha\frac{1}{m}\sum_{i=1}^{m}\|\mathbf{T_{i}}\|_{1}}_{\mathrm{\mathcal{L}_{1}}}+\underbrace{\frac{1}{m}\sum_{i=1}^{m}(\mathbf{T^{\prime\prime}_{i}}-\mathbf{T_{i}})^{2}}_{\mathrm{\mathcal{L}_{2}}}+\underbrace{\frac{1}{m}\sum_{i=1}^{m}(\mathbf{T^{\prime\prime}_{i}}-\mathbf{T^{\prime}_{i}})^{2}}_{\mathrm{\mathcal{L}_{3}}}$
(7)

where $\alpha=0.2$ is the coefficient for the direct sparsity constraint term.

##  3 Results

###  3.1 Validation of the sparsity of transformed TMs

We validate the effectiveness of the basis transformation by assessing the
change in sparsity of the transformed representation of TMs, which can be
measured by participation ratio or by the latent space size of an autoencoder
(see Appendix Eqn. 9). The participation ratio, $p$, taking values between 0
to 1, quantifies the fraction of ‘occupied’ matrix elements (see Appendix Eqn.
8). Thus a lower participation ratio represents a sparser distribution of
power within the elements of the matrix. Figure 3 shows three examples of the
transformed representation of randomly-generated forward TMs,
$\mathbf{T_{f}}\in\mathbb{C}^{78\times 78}$ (which appear relatively sparse in
their initial form, with a mean participation ratio of $p=0.06$), randomly-
generated round-trip TMs, $\mathbf{T_{r}}\in\mathbb{C}^{78\times 78}$ (which
appear dense in their initial form, with mean $p=0.46$) and physically-modeled
TMs $\mathbf{T_{e}}\in\mathbb{C}^{78\times 78}$ (most sparse in their original
form, with mean $p=0.03$). We test five different models, namely linear
similarity transformation, CNN, FCNN, single self-attention layer and self-
attention based-FCNN model.

For the dense and relatively sparse datasets, the linear model shows limited
capability in transforming a group of matrices into a more sparse
representation, with transformed mean participation ratios similar to the
values of the original matrices.

Compared to the FCNN, the CNN presents a significantly larger mean
participation ratio, especially for dense matrices where the $p$ value is two
times larger than that achieved using the FCNN. This demonstrates that CNNs,
with their inbuilt bias for local correclations, have limited capability for
handling non-locally correlated matrices (such as fibre TMs). On the other
hand, the simple self-attention layer, which learns position-wise weights,
achieves a better performing non-linear transformation than the FCNN, with a
further reduced mean participation ratio. This is because the attention
mechanism selectively focuses on relevant parts of the matrices by computing
alignment scores, allowing the model to dynamically prioritize information
based on the input itself, while fully connected layers uniformly apply the
same transformation to the entire input data without any mechanism to focus
more on specific parts.

To further enhance the data compressibility, the self-attention-based FCNN
model integrates the attention mechanism into the FCNN framework and further
decreases the mean participation ratio to $p=0.03$ for relatively sparse
matrices and $p=0.11$ for dense matrices.

Figure 3: TM transformations using different models. The three coloured boxes
show transformations on different data types. Within each box, the left hand
column shows the TM transformed into the sparse representation, the middle
column shows the TM transformed back into its original form and the right hand
column shows the residual error in the TM after having been compressed and
restored. Each row shows data transformed using a different model: (a) Linear
similarity transformation based model, (b) CNN, (c) FCNN, (d) self-attention
only, (e) self-attention-based FCNN model. Metrics are represented by mean
$\pm$ standard deviation.

###  3.2 Validation of the invertibility of transformed TMs

We next examine the invertibility of the transformed TMs, which reflects the
model’s ability to preserve the information within the original matrices.
Figure 3 illustrates three examples of TMs reconstructed from the transformed
TMs. The linear model provides a lossless invertible process, with a basis
matrix identified by the pre-trained model. Among the non-linear models, the
FCNN, the simple self-attention layer, and the self-attention-based FCNN model
are all capable of recovering the TMs to the original basis with an error of
$\leq 8\%$ for sparse matrices and $\leq 14\%$ for dense matrices. The CNN
presents fair reconstruction for sparse matrices, with a slightly larger error
of 12%, but exhibits a large mean error of 32% for dense matrices, indicating
a limited ability for reconstruction.

###  3.3 Compatibility of physically-modeled TMs

To further test if our methods are compatible with physically modeled TMs, we
employed a group of physically modeled perturbed TMs to train the self-
attention-based FCNN model. Figure 3 shows the representation of TMs
transformed from the original basis using the different models, with the mean
participation ratio ranging from $p=0.05$ (linear model) to $p=0.01$ (self
attention-based FCNN). Additionally, this approach allows for the
reconstruction of TMs from the transformed basis with the lowest mean error of
5%.

##  4 Discussion

We haveintroduced a new approach to learn critical features of ensembles of
fibre TMs, approaching a representation of these TMs that is maximally sparse
and basis invariant. Our methods were applied to diverse TM datasets,
including randomly-generated and physically-modeled matrices under dynamic
conditions, with various sparsity properties. We tested five models, and Table
1 summarizes the comparative performance, expressed in terms of the matrix
compressibility, defined by participation ratio $p$ (Eqn. 8) and the size of
latent space from an autoencoder (Eqn. 9), and the invertibility measured by
the reconstruction error (Eqn. 10).

It is evident that non-linear models generally outperform the linear model. In
general, we expect the participation ratio to scale with the dimension of the
TMs, where (ignoring any correlation between diagonal elements) we expect the
lowest possible participation ratio is given by a perfectly diagonal TM,
yielding $p=1/N$, where $N$ is the matrix dimension. The self-attention-based
FCNN model shows superior performance in achieving lower mean errors and lower
participation ratios compared to other models. In the case of physically
modelled perturbed fibres, $p=0.01\pm 0.017$, which is similar to the
participation ratio of a perfectly diagonal TM ($p=1/78=0.013$). This means
the self-attention mechanism has the ability to dynamically learn the relevant
features across varying TM characteristics.

Table 1: Comparison of transformation performance across diverse datasets
using different models

| Data Properties | Metric | Linear | CNN | FCNN |  | Self-Attention
---
only
| Self-attention
---
-based FCNN
|  | $p$ | 0.07$\pm$0.010 | 0.13$\pm$0.041 | 0.04$\pm$0.021 | 0.04$\pm$0.018 | 0.03$\pm$0.009
|  | LS ratio | NA | 14.79% | 5.33% | 5.33% | 4.21%
Dataset 1: Randomly-generated forward TMs | Less sparse | mean error | 0 | 0.12$\pm$0.03 | 0.07$\pm$0.04 | 0.07$\pm$0.03 | 0.06$\pm$0.02
|  | $p$ | 0.39$\pm$0.141 | 0.36$\pm$0.183 | 0.15$\pm$0.040 | 0.14$\pm$0.035 | 0.11$\pm$0.016
|  | LS ratio | NA | >67.32% | 16.83% | 16.83% | 14.79%
Dataset 2: Randomly-generated round-trip TMs | Dense | mean error | 0 | 0.32$\pm$0.21 | 0.12$\pm$0.09 | 0.14$\pm$0.07 | 0.08$\pm$0.02
|  | $p$ | 0.05$\pm$0.021 | 0.04$\pm$0.020 | 0.03$\pm$0.014 | 0.02$\pm$0.019 | 0.01$\pm$0.017
|  | LS ratio | NA | 7.25% | 4.21% | 4.21% | 3.22%
Dataset 3: Physical-modeled perturbed TMs | Sparse | mean error | 0 | 0.07$\pm$0.02 | 0.06$\pm$0.02 | 0.05$\pm$0.01 | 0.05$\pm$0.01

To intuit how these models work, we can draw a comparison with pre-existing
linear algebra approaches. TM diagonalisation (i.e. maximal sparsity for a
full-rank matrix) is achievable by measuring the TM for any specific fibre,
extracting the eigenvectors and using these as the basis. We consider this a
form of ‘overfitting’ as each solution is entirely specific to the particular
TM. By contrast, a fixed linear basis transformation that is optimised so as
to result in maximally sparse representations across a wide range of TMs may
provide some improvement over the worst case fully dense basis, but we
consider this a form of ‘underfitting’. A third way, which we propose here, is
to develop a non-linear basis transformation that is input-dependent but is
constrained or regularized to utilize few internal parameters. We show how
this approach enables a nonlinear basis transformation of TMs that models
perturbations in an efficient way. Intuitively, we expect such perturbations
to be highly compressible since they may be described by a small number of
physical parameters: curvature as a function of position along the fibre
(which is relatively simple due to the rigidity of glass), degree of twist,
minor manufacturing variations in a relatively smooth refractive index
profile, a small number of optical defects (which accounts for the low-loss of
optical fibres) and temperature. This means we expect an ensemble of fibre TMs
to represent a minute subset of the set of all possible TMs – and the task of
our model is to identify the correlations shared by TMs in this space.

To understand the effectiveness of self-attention mechanisms, we examine their
widespread use in language modelling. Many large language models effectively
take sequences of words represented as vectors and transform these vectors to
a more representative basis that better encodes semantics. By considering
matrices to be analogous to blocks of text (e.g. sentences) and columns
analogous to word embedding vectors represented in the measurement basis, we
can use self-attention layers to learn input-dependent basis transformations
that better represents TMs under perturbations. Previous approaches to optical
fibre imaging have relied on CNNs to model a static TM. However, because of
the arbitrary choice of basis when modelling over a range of TMs, CNNs are not
suitable as they contain a strong inductive bias that spatially adjacent
elements are correlated, which is not at all the case here. Rather, the
attention mechanism allows correlations between elements to be learned.

###  4.1 Limitations

Several improvements should be considered for further refining this approach.
First, our methods require a large amount of training data which could be
difficult to obtain if using real experimentally measured TMs. This can be
addressed by augmenting experimental data with a large amount of simulated
fiber data. Approaches that use generative networks for data augmentation and
domain transfer [51, 52] can enable convergence using relatively small
experimental datasets, or else help to reliably generate further training
data. Second, the training process is highly memory-intensive due to the large
sizes of TMs typically used in imaging applications. Large-language models
(LLMs) can have a word-embedding dimension of $\geq$10,000 with > 5000 tokens
sequence lengths. This provides an idea of the potential for scaling up our
self-attention-based model: complex-valued matrices of size $\geq$ 3000
$\times$ 3000 should be comfortably feasible using these frameworks. While the
use of self-attention has marginally reduced memory consumption, self-
attention networks grow with the square of the embedding dimension. Since our
transmission matrices have a number of elements equal to $x^{2}\times x^{2}$
for input fields (i.e./ images) containing $x\times x$ pixels, the self-
attention architecture grows as as the fourth power of the image dimension
i.e. $x^{4}$. This means that processing very high resolution images may be
challenging. Third, to generate random matrices used for training, a limited
fiber bandwidth was assumed to avoid degeneracies arising from matrix
logarithms. In future the spectral correlations of fibre TMs could be
exploited by AI models by using a non-linear model of the fiber TM over a much
broader wavelength range [2].

If used in a healthcare imaging setting, some consideration needs to be given
to spurious imaging results or the potential for adversarial attacks. However,
this can be mitigated somewhat by using predicted transmission matrices to
send pilot signals down a fiber and verify that the observed behaviour of the
physical system matches that of the model’s prediction. Such physical
verification adds an additional level of safety.

In the future, this fibre TM representation architecture may facilitate the
development of flexible hair-thin optical imaging systems that adapt to fibre
changes in real time, by enabling continuous fibre matrix estimation with a
minimal re-characterisation measurements that are tailored to minimise
uncertainty in the latent space. Our work may also allow practical modelling
of complex effects such as wavelength dispersion and nonlinear
characterisation, and the exploration of generative models for predicting full
TMs given partial information. This could in turn enable a wide range of
broadband imaging modalities to be implemented through optical fibres.

Finally, we anticipate that our self-attention-based neural network method
will lead to the new AI models that are invariant or equivariant to arbitrary
modal representation bases.

## DATA AVAILABILITY

The data presented in this study are available from the following source: [DOI
to be inserted later].

## CODE AVAILABILITY

The code for this study is available from the following source: [DOI to be
inserted later].

## ACKNOWLEDGEMENT

DBP thanks the European Research Council for financial support (PhotUntangle
804626). RJK was supported by an EPSRC doctoral training partnership grant
EP/T518049/1. GSDG acknowledges support from a UKRI Future Leaders Fellowship
(MR/T041951/1) and an EPSRC Ph.D. studentship.

## Appendix - Evaluation metric

To assess the effectiveness of transformation of the TMs,
$\mathbf{T_{1...m}}\in\mathbb{C}^{N\times N}$ (where here $N=78$), we measure
the mean participation ratio [53] of the transformed TMs and define a metric,
$p$, to quantify the fraction of occupied matrix elements (i.e. representing
the sparsity), where $\mathbf{T_{j}}$ represents $j^{\mathrm{th}}$ element of
each TM:

$p=\frac{(\sum_{i=1}^{m}\sum_{j=1}^{N\times
N}|\mathbf{T_{j}}|^{4})^{-1}}{N^{2}}$ (8)

Additionally, we construct an autoencoder model to downsample the transformed
TM into a low-dimensional LS1 which is expected to contain the important
features that can reconstruct the transformed TM. The compressibility (i.e.
sparsity) of the original TMs, can be determined by:

$\mathrm{LS_{Ratio}}=\frac{1}{m}\sum_{i=1}^{m}\frac{\texttt{Size of
ls1}_{i}}{2N\times 2N}$ (9)

To further evaluate the invertibility of the transformed TMs, we construct a
separate autoencoder model, where the decoder is expected to reconstruct the
original TMs from LS1. We calculate mean squared error across the test dataset
(mean error) between the reconstructed TMs and target TMs:

$\mathrm{mean}\,\mathrm{error}=\frac{1}{m}\sum_{i=1}^{m}\sum(\mathbf{T^{\prime\prime}_{i}}-\mathbf{T_{i}})^{2}$
(10)

## References

  * [1] Demetri Psaltis and Christophe Moser.  Imaging with multimode fibers.  Optics and Photonics News, 27(1):24–31, 2016.
  * [2] Szu-Yu Lee, Vicente J Parot, Brett E Bouma, and Martin Villiger.  Efficient dispersion modeling in optical multimode fiber.  Light: Science & Applications, 12(1):31, 2023.
  * [3] Shuhui Li, Charles Saunders, Daniel J Lum, John Murray-Bruce, Vivek K Goyal, Tomáš Čižmár, and David B Phillips.  Compressively sampling the optical transmission matrix of a multimode fibre.  Light: science & applications, 10(1):88, 2021.
  * [4] Martin Plöschner, Tomáš Tyc, and Tomáš Čižmár.  Seeing through chaos in multimode fibres.  Nature Photonics, 9(8):529–535, 2015.
  * [5] Shengfu Cheng, Tianting Zhong, Chi Man Woo, Qi Zhao, Hui Hui, and Puxiang Lai.  Long-distance pattern projection through an unfixed multimode fiber with natural evolution strategy-based wavefront shaping.  Optics Express, 30(18):32565–32576, 2022.
  * [6] Robert Kuschmierz, Elias Scharf, Nektarios Koukourakis, and Jürgen W Czarske.  Self-calibration of lensless holographic endoscope using programmable guide stars.  Optics letters, 43(12):2997–3000, 2018.
  * [7] Nicolino Stasio, Donald B Conkey, Christophe Moser, and Demetri Psaltis.  Light control in a multicore fiber using the memory effect.  Optics express, 23(23):30532–30544, 2015.
  * [8] Uri Weiss and Ori Katz.  Two-photon lensless micro-endoscopy with in-situ wavefront correction.  Optics Express, 26(22):28808–28817, 2018.
  * [9] Chi Man Woo, Huanhao Li, Qi Zhao, and Puxiang Lai.  Dynamic mutation enhanced particle swarm optimization for optical wavefront shaping.  Optics Express, 29(12):18420–18426, 2021.
  * [10] Chi Man Woo, Qi Zhao, Tianting Zhong, Huanhao Li, Zhipeng Yu, and Puxiang Lai.  Optimal efficiency of focusing diffused light through scattering media with iterative wavefront shaping.  APL Photonics, 7(4), 2022.
  * [11] Qi Zhao, Chi Man Woo, Huanhao Li, Tianting Zhong, Zhipeng Yu, and Puxiang Lai.  Parameter-free optimization algorithm for iterative wavefront shaping.  Optics Letters, 46(12):2880–2883, 2021.
  * [12] Antonio M Caravaca-Aguirre, Adrien Carron, Sylvain Mezil, Irène Wang, and Emmanuel Bossy.  Optical memory effect in square multimode fibers.  Optics Letters, 46(19):4924–4927, 2021.
  * [13] Nicolino Stasio, Christophe Moser, and Demetri Psaltis.  Calibration-free imaging through a multicore fiber using speckle scanning microscopy.  Optics letters, 41(13):3078–3081, 2016.
  * [14] Weiru Fan, Ziyang Chen, Vladislav V Yakovlev, and Jixiong Pu.  High-fidelity image reconstruction through multimode fiber via polarization-enhanced parametric speckle imaging.  Laser & Photonics Reviews, 15(5):2000376, 2021.
  * [15] Amir Porat, Esben Ravn Andresen, Hervé Rigneault, Dan Oron, Sylvain Gigan, and Ori Katz.  Widefield lensless imaging through a fiber bundle via speckle correlations.  Optics express, 24(15):16835–16855, 2016.
  * [16] Antonio M Caravaca-Aguirre, Sakshi Singh, Simon Labouesse, Michael V Baratta, Rafael Piestun, and Emmanuel Bossy.  Hybrid photoacoustic-fluorescence microendoscopy through a multimode fiber using speckle illumination.  Apl Photonics, 4(9), 2019.
  * [17] Antonio M Caravaca-Aguirre, Florian Poisson, Dorian Bouchet, Nicolino Stasio, Philippe Moreau, Irene Wang, Edward Zhang, Paul Beard, Claire Prada, Christophe Moser, et al.  Single-pixel photoacoustic microscopy with speckle illumination.  Intelligent Computing, 2:0011, 2023.
  * [18] Lyubov V Amitonova and Johannes F De Boer.  Compressive imaging through a multimode fiber.  Optics letters, 43(21):5427–5430, 2018.
  * [19] Liang Deng, Joseph D Yan, Daniel S Elson, and Lei Su.  Characterization of an imaging multimode optical fiber using a digital micro-mirror device based single-beam system.  Optics express, 26(14):18436–18447, 2018.
  * [20] Moussa N’Gom, Theodore B Norris, Eric Michielssen, and Raj Rao Nadakuditi.  Mode control in a multimode fiber through acquiring its transmission matrix from a reference-less optical system.  Optics letters, 43(3):419–422, 2018.
  * [21] Guoqiang Huang, Daixuan Wu, Jiawei Luo, Yin Huang, and Yuecheng Shen.  Retrieving the optical transmission matrix of a multimode fiber using the extended kalman filter.  Optics express, 28(7):9487–9500, 2020.
  * [22] Shengfu Cheng, Xuyu Zhang, Tianting Zhong, Huanhao Li, Haoran Li, Lei Gong, Honglin Liu, and Puxiang Lai.  Nonconvex optimization for optimum retrieval of the transmission matrix of a multimode fiber.  Advanced Photonics Nexus, 2(6):066005–066005, 2023.
  * [23] Ang Li, Wenxuan Liang, Honghua Guan, Yung-Tian A Gau, Dwight E Bergles, and Xingde Li.  Focus scanning with feedback-control for fiber-optic nonlinear endomicroscopy.  Biomedical optics express, 8(5):2519–2527, 2017.
  * [24] Shuhui Li, Simon AR Horsley, Tomáš Tyc, Tomáš Čižmár, and David B Phillips.  Memory effect assisted imaging through multimode optical fibres.  Nature Communications, 12(1):3751, 2021.
  * [25] Salma Farahi, David Ziegler, Ioannis N Papadopoulos, Demetri Psaltis, and Christophe Moser.  Dynamic bending compensation while focusing through a multimode fiber.  Optics express, 21(19):22504–22514, 2013.
  * [26] Zhong Wen, Zhenyu Dong, Qilin Deng, Chenlei Pang, Clemens F Kaminski, Xiaorong Xu, Huihui Yan, Liqiang Wang, Songguo Liu, Jianbin Tang, et al.  Single multimode fibre for in vivo light-field-encoded endoscopic imaging.  Nature Photonics, 17(8):679–687, 2023.
  * [27] Ruo Yu Gu, Reza Nasiri Mahalati, and Joseph M Kahn.  Design of flexible multi-mode fiber endoscope.  Optics express, 23(21):26905–26918, 2015.
  * [28] Haoshuo Chen, Nicolas K Fontaine, Roland Ryf, David T Neilson, and Peter Winzer.  Remote spatio-temporal focusing over multimode fiber enabled by single-ended channel estimation.  IEEE Journal of Selected Topics in Quantum Electronics, 26(4):1–9, 2020.
  * [29] George SD Gordon, Milana Gataric, Alberto Gil CP Ramos, Ralf Mouthaan, Calum Williams, Jonghee Yoon, Timothy D Wilkinson, and Sarah E Bohndiek.  Characterizing optical fiber transmission matrices using metasurface reflector stacks for lensless imaging without distal access.  Physical review X, 9(4):041050, 2019.
  * [30] Maxime W Matthès, Yaron Bromberg, Julien de Rosny, and Sébastien M Popoff.  Learning and avoiding disorder in multimode fibers.  Physical Review X, 11(2):021060, 2021.
  * [31] Babak Rahmani, Damien Loterie, Georgia Konstantinou, Demetri Psaltis, and Christophe Moser.  Multimode optical fiber transmission with a deep learning network.  Light: science & applications, 7(1):69, 2018.
  * [32] Piergiorgio Caramazza, Oisín Moran, Roderick Murray-Smith, and Daniele Faccio.  Transmission of natural scene images through a multimode fibre.  Nature communications, 10(1):2029, 2019.
  * [33] Pengfei Fan, Tianrui Zhao, and Lei Su.  Deep learning the high variability and randomness inside multimode fibers.  Optics express, 27(15):20241–20258, 2019.
  * [34] Pengfei Fan, Michael Ruddlesden, Yufei Wang, Luming Zhao, Chao Lu, and Lei Su.  Learning enabled continuous transmission of spatially distributed information through multimode fibers.  Laser & Photonics Reviews, 15(4):2000348, 2021.
  * [35] Navid Borhani, Eirini Kakkava, Christophe Moser, and Demetri Psaltis.  Learning to see through multimode fibers.  Optica, 5(8):960–966, 2018.
  * [36] Shachar Resisi, Sebastien M Popoff, and Yaron Bromberg.  Image transmission through a dynamically perturbed multimode fiber by deep learning.  Laser & Photonics Reviews, 15(10):2000553, 2021.
  * [37] Oisín Moran, Piergiorgio Caramazza, Daniele Faccio, and Roderick Murray-Smith.  Deep, complex, invertible networks for inversion of transmission effects in multimode optical fibres.  Advances in neural information processing systems, 31, 2018.
  * [38] Guohua Wu, Yong Sun, Longfei Yin, Zhixiong Song, and Wenting Yu.  High-definition image transmission through dynamically perturbed multimode fiber by a self-attention based neural network.  Optics Letters, 48(10):2764–2767, 2023.
  * [39] Yijie Zheng, Terry Wright, Zhong Wen, Qing Yang, and George SD Gordon.  Single-ended recovery of optical fiber transmission matrices using neural networks.  Communications Physics, 6(1):306, 2023.
  * [40] Grégoire Delétang, Anian Ruoss, Paul-Ambroise Duquenne, Elliot Catt, Tim Genewein, Christopher Mattern, Jordi Grau-Moya, Li Kevin Wenliang, Matthew Aitchison, Laurent Orseau, et al.  Language modeling is compression.  arXiv preprint arXiv:2309.10668, 2023.
  * [41] Gongbo Tang, Mathias Müller, Annette Rios, and Rico Sennrich.  Why self-attention? a targeted evaluation of neural machine translation architectures.  arXiv preprint arXiv:1808.08946, 2018.
  * [42] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani.  Self-attention with relative position representations.  arXiv preprint arXiv:1803.02155, 2018.
  * [43] Yang Liu and Mirella Lapata.  Text summarization with pretrained encoders.  arXiv preprint arXiv:1908.08345, 2019.
  * [44] Song Xu, Haoran Li, Peng Yuan, Youzheng Wu, Xiaodong He, and Bowen Zhou.  Self-attention guided copy mechanism for abstractive summarization.  In Proceedings of the 58th annual meeting of the association for computational linguistics, pages 1355–1362, 2020.
  * [45] Walaa Medhat, Ahmed Hassan, and Hoda Korashy.  Sentiment analysis algorithms and applications: A survey.  Ain Shams engineering journal, 5(4):1093–1113, 2014.
  * [46] Praneeth Chakravarthula, Jipeng Sun, Xiao Li, Chenyang Lei, Gene Chou, Mario Bijelic, Johannes Froesch, Arka Majumdar, and Felix Heide.  Thin on-sensor nanophotonic array cameras.  ACM Transactions on Graphics (TOG), 42(6):1–18, 2023.
  * [47] Stefan Rotter and Sylvain Gigan.  Light fields in complex media: Mesoscopic scattering meets wave control.  Reviews of Modern Physics, 89(1):015005, 2017.
  * [48] Thomas Chaigne, Ori Katz, A Claude Boccara, Mathias Fink, Emmanuel Bossy, and Sylvain Gigan.  Controlling light in scattering media non-invasively using the photoacoustic transmission matrix.  Nature Photonics, 8(1):58–64, 2014.
  * [49] Ulrike Hahn, Nick Chater, and Lucy B Richardson.  Similarity as transformation.  Cognition, 87(1):1–32, 2003.
  * [50] Gerald N Hile and Pertti Lounesto.  Matrix representations of clifford algebras.  Linear algebra and its applications, 128:51–63, 1990.
  * [51] Sebastian J Wirkert, Anant S Vemuri, Hannes G Kenngott, Sara Moccia, Michael Götz, Benjamin FB Mayer, Klaus H Maier-Hein, Daniel S Elson, and Lena Maier-Hein.  Physiological parameter estimation from multispectral images unleashed.  In Medical Image Computing and Computer Assisted Intervention- MICCAI 2017: 20th International Conference, Quebec City, QC, Canada, September 11-13, 2017, Proceedings, Part III 20, pages 134–141. Springer, 2017.
  * [52] Ahmed Osman, Jane Crowley, and GS D Gordon.  Training generative adversarial networks for optical property mapping using synthetic image data.  Biomedical Optics Express, 13(10):5171–5186, 2022.
  * [53] Bart A Van Tiggelen.  Localization of waves.  Diffuse waves in complex media, pages 1–60, 1999.