🤖 machine learning

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

This paper introduces a self-attention-based framework that dynamically transforms multimode optical fibre transmission matrices into compact, low-dimensional latent spaces, effectively addressing the challenges of dynamic environmental changes and non-linearities to achieve high-fidelity image reconstruction with minimal error.

Original authors: Yijie Zheng, Robert J. Kilpatrick, David B. Phillips, George S. D. Gordon

Published 2026-06-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yijie Zheng, Robert J. Kilpatrick, David B. Phillips, George S. D. Gordon

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Scrambled Message" Problem

Imagine you have a very thin, hair-like glass thread (an optical fiber) that you want to use as a camera to look inside the human body. Because the thread is so thin, it can slip into tiny, hard-to-reach places like deep blood vessels or the brain.

However, there's a catch. As light travels through this thread, it gets scrambled. Think of it like taking a clear photo, putting it in a blender, and then trying to guess what the original picture looked like just by looking at the smoothie.

In the past, scientists used a "static map" (a pre-calculated list of rules) to unscramble the image. But this map only works if the fiber stays perfectly still. In the real world, the fiber bends, twists, and changes temperature, which completely messes up the map. It's like trying to use a paper map of a city that is constantly being reshaped by earthquakes.

The Solution: A "Smart Translator"

The authors of this paper built a new type of computer brain (a neural network) that acts like a smart translator. Instead of trying to memorize a fixed map, this translator learns how to instantly reorganize the scrambled data into a neat, compact format that is easy to understand, even when the fiber is moving.

Here is how they did it, broken down into three key ideas:

1. The "Long-Distance Phone Call" (Self-Attention)

Most computer brains used for images (called CNNs) work like a person looking at a photo and only paying attention to the pixels right next to each other. They assume neighbors are related.

But in these fiber threads, a pixel on the far left might be mathematically connected to a pixel on the far right due to how the light bounces around inside the glass. It's like a long-distance phone call where the person at the start of the line is talking to the person at the very end, skipping everyone in the middle.

The authors used a technology called Self-Attention (the same tech behind modern AI chatbots). Instead of just looking at neighbors, this system lets every part of the data "talk" to every other part instantly. This allows the computer to see the whole picture of how the light is scrambled, no matter how far apart the connections are.

2. The "Suitcase Packing" (Latent Space Compression)

The scrambled data from the fiber is huge and messy, like a suitcase stuffed with clothes thrown in randomly.

Old way: You try to carry the whole messy suitcase. It's heavy and hard to manage.
New way: The authors' model acts like a master packer. It takes that messy suitcase and folds everything into a tiny, neat, compact cube (a "latent space").

This "compact cube" is a simplified version of the data that keeps all the important information but throws away the clutter. The paper shows that they can shrink the data down to less than 10% of its original size (making it very "sparse") without losing the ability to unpack it later.

3. The "Universal Adapter" (Basis Invariance)

One of the biggest headaches in this field is that scientists measure light in different "languages" or "bases" (like measuring in pixels, or in waves, or in patterns). Usually, a computer model trained in one language fails if you switch to another.

The authors' model is like a universal adapter. They tested it by feeding it data in three different "languages" (LP, Fourier, and Hadamard bases). The model successfully translated all of them into the same neat, compact format. This means the model understands the physics of the fiber, not just the specific way the data was written down.

How They Tested It

The team didn't just guess; they ran rigorous tests:

Simulated Fibers: They created thousands of fake fibers on a computer, bending and twisting them to see if the model could handle the chaos.
Real Experiments: They used actual glass fibers in a lab, bending them physically, and the model still worked.
The "Unscramble" Test: They compressed the data and then tried to rebuild the original image. The model was able to reconstruct the original data with less than 10% error, proving it didn't throw away anything important.

The Bottom Line

The paper claims that by using Self-Attention (which looks at long-range connections) and Compression (folding data into a small space), they have created a system that can handle the messy, moving reality of optical fibers much better than previous methods.

It's like upgrading from a static paper map to a GPS that instantly recalculates the route every time the road changes, ensuring you can still find your way even if the terrain is shifting beneath your feet. This makes the idea of using hair-thin fibers for clear, real-time medical imaging much more possible.

Technical Summary: Self-Attention Networks for Compact Latent Space Modelling of Dynamic Optical Fibre Matrices

Problem Statement

Multimode optical fibres (MMFs) offer a pathway for ultra-thin, minimally invasive medical endoscopy capable of sub-cellular imaging deep within the body. However, the transmission of light through these fibres inherently scrambles images due to modal coupling. While conventional methods utilize pre-calibrated linear transmission matrices (TMs) to correct this scrambling, they fail in dynamic real-world scenarios. Physical perturbations—such as fibre movement, temperature variations, and bending—alter the TM continuously. Furthermore, obtaining the input-output pairs required for recalibration is often impractical in clinical settings where the fibre tip is inaccessible.

Existing neural network approaches, particularly Convolutional Neural Networks (CNNs), struggle with this task. CNNs rely on the inductive bias of local pixel correlations, an assumption invalidated by the long-range correlations present in fibre TMs. These correlations arise from both physical non-local modal coupling and the mathematical necessity of flattening multidimensional optical fields into matrix representations. Consequently, conventional architectures require prohibitively large parameter counts, lack robustness to changes in measurement bases (e.g., pixel vs. Fourier vs. Hadamard), and fail to capture the complex dynamics of perturbed fibres efficiently.

Methodology

The authors propose a novel approach using self-attention layers to dynamically transform varying fibre TMs into compact, low-dimensional latent-space representations. The core philosophy follows "compression is learning," where the network learns to map large, complex TMs into a sparse latent space suitable for downstream tasks like image reconstruction or bending estimation.

Data Generation and Datasets

The study validates the approach using four diverse datasets of $78 \times 78$ complex-valued TMs:

Randomly-generated forward TMs: Simulated to exhibit high sparsity and specific power intensity distributions typical of real fibres.
Randomly-generated round-trip TMs: Computed as $T_r = T_f^T R T_f$ , where $R$ is a scattering reflector matrix. These represent dense matrices relevant to single-ended imaging.
Physically-modelled TMs: Generated using a segmented fibre model simulating bends and curvatures over lengths of 0.6–1.2 m.
Experimental TMs: Measured using an off-axis holography setup with a digital micro-mirror device (DMD) and a CCD camera, covering 164 different bending conformations.

Model Architecture

The authors compare several architectures for basis transformation:

Linear Similarity Transformation: A differentiable layer learning a basis matrix $B$ to transform $T' = B^{-1}TB$ .
Convolutional Neural Networks (CNN): Standard 2D convolutions with LeakyReLU activations.
Fully Connected Neural Networks (FCNN): Dense layers with LeakyReLU activations.
Self-Attention Only: A single self-attention layer.
Self-Attention + FCNN: The proposed hybrid model where a multi-head self-attention mechanism precedes two fully connected layers.

Key Mechanism: The input TMs are treated as sequences of modal channels. The self-attention mechanism computes Query ( $Q$ ), Key ( $K$ ), and Value ( $V$ ) matrices to capture long-range pairwise dependencies between elements, allowing the model to learn non-local interactions without the spatial locality bias of CNNs.

Training Constraints

To ensure the transformation is physically meaningful and reversible, the models are trained with a composite loss function:

Sparsity Constraint ( $L_1$ ): Minimizes the $L_1$ norm of the transformed matrix to encourage a sparse latent representation.
Invertibility Constraint ( $L_2$ ): An autoencoder reconstructs the original TM from the transformed representation to ensure information preservation.
Sparsity Reconstruction Constraint ( $L_3$ ): A second autoencoder ensures the compressed latent space contains sufficient information to reconstruct the transformed sparse matrix.
Complex Number Handling: Real and imaginary parts are separated, and a custom loss function invariant to global phase shifts is used.

Key Contributions

Introduction of Self-Attention for TM Modelling: The paper demonstrates that self-attention mechanisms are uniquely suited to model optical fibres because they capture the long-range, non-local correlations inherent in transmission matrices, which CNNs miss.
Compact Latent Space Representation: The proposed method transforms dense and sparse TMs into highly sparse representations (participation ratios between 0.01 and 0.11), significantly reducing dimensionality while preserving invertibility.
Basis Invariance: The learned transformation maintains structural consistency across different measurement bases (LP/PIM, Fourier, Hadamard), addressing a critical limitation of previous models that overfit to specific bases.
Robustness to Dynamic Perturbations: The model demonstrates tolerance to dynamic transmission matrix perturbations, effectively handling scenarios where the fibre state changes during measurement (simulated by swapping columns in the TM).

Results

Sparsity and Efficiency: The Self-Attention + FCNN model achieved the lowest participation ratios ( $p \approx 0.03$ for sparse matrices and $0.11$ for dense matrices) and the smallest latent space ratios (down to 3.22% for dense matrices), outperforming CNNs, FCNNs, and linear models.
Reconstruction Accuracy: The model successfully reconstructed original TMs with errors of $\le 8\%$ for sparse matrices and $\le 14\%$ for dense matrices. In contrast, CNNs exhibited large errors (up to 32%) for dense matrices.
Cross-Basis Consistency: Cosine similarity between transformed matrices of the same physical state but different bases was significantly higher (approx. 0.80) for the self-attention model compared to the original matrices (approx. 0.65), indicating the model learns the underlying physical structure rather than basis-specific artifacts.
Generalization: A model trained on physically simulated perturbed TMs successfully generalized to experimentally measured TMs without retraining, reducing the participation ratio from 0.37 to 0.12.
Perturbation Robustness: The model maintained low reconstruction errors (<22%) under dynamic perturbations (column swapping) up to rates of 13%, whereas CNNs failed to achieve non-linear transformation effectively under these conditions.

Significance and Claims

The authors claim that this approach provides a flexible, memory-efficient, and robust framework for modelling dynamic fibre TMs. By converting complex, high-dimensional matrices into compact latent spaces, the method facilitates single-ended fibre imaging, potentially enabling the deployment of needle-sized imaging devices.

The paper emphasizes that the success of self-attention lies in its ability to treat transmission matrices as structured sequences, learning input-dependent transformations that adapt to perturbations. Unlike fixed linear bases or locality-biased CNNs, this data-dependent approach captures the non-local modal coupling essential for accurate fibre modelling.

The authors modestly note that while the model captures recurring structures, individual latent variables are not yet explicitly mapped to physical coordinates (e.g., specific bend radii). However, the work lays the foundation for future real-time distortion correction, missing data imputation, and the application of similar matrix compression techniques to other scattering media. The paper acknowledges current limitations, such as the quadratic scaling of attention complexity with matrix size and the use of real-valued optimization for complex matrices, suggesting future work in state-space models (e.g., Mamba) and fully complex-valued training.