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

TL;DR

The paper addresses dynamic, nonlinear transmission matrices (TMs) in multimode fibre imaging, where static linear models fail under perturbations. It introduces self-attention-based basis transformations to map TM ensembles into compact, sparsity-promoting representations while enforcing invertibility via autoencoder constraints. Three TM datasets (forward, round-trip, and physically modeled) are used to compare linear similarity transforms and non-linear architectures (CNN, FCNN, self-attention, and self-attention+FCNN), with performance measured by participation ratio $p$ and reconstruction error. Results show that self-attention-based models dramatically reduce $p$ (to about $0.01$–$0.11$ in various cases) while maintaining low reconstruction errors, enabling effective dynamic TM modelling from single-ended measurements and paving the way for real-time fibre imaging and broader applications in scattering media, wavelength dispersion, and nonlinear effects.

Abstract

Multimode optical fibres are hair-thin strands of glass that efficiently transport light. They promise next-generation medical endoscopes that provide unprecedented sub-cellular image resolution deep inside the body. However, confining light to such 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 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 varying fibre matrices to a basis that admits compact, low-dimensional representations suitable for further processing. We demonstrate the effectiveness of this approach on diverse fibre matrix datasets. 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 < 10% reconstruction error, demonstrating the invertibility.

Figures (3)

• 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.
• 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.
• 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.