Utilising a learned forward operator in the inverse problem of photoacoustic tomography

This paper demonstrates that a Fourier neural operator can serve as an accurate and computationally efficient learned forward operator for solving the inverse problem in photoacoustic tomography via gradient-based optimization, outperforming conventional pseudospectral $k$-space methods in efficiency while maintaining high accuracy.

The Gist

The Big Picture: "Seeing" Inside Without Cutting

Imagine you want to know what's inside a sealed, opaque box. You can't open it, but you can shine a flashlight on it. When the light hits something inside, it heats up slightly and creates a tiny "pop" (a sound wave). If you listen carefully to the sounds coming out of the box, you can try to figure out what's inside.

This is Photoacoustic Tomography (PAT). It's a medical imaging technique that uses light to create sound, and then listens to that sound to build a picture of what's inside the body (like blood vessels or tumors).

The Problem: The "Slow Calculator"

To turn those sounds back into a clear picture, scientists have to solve a massive math puzzle called an inverse problem. They have to work backward: If I heard this sound, what did the object inside look like?

To solve this, they need a Forward Operator. Think of this as a "Simulation Machine."

• The Old Way: The traditional simulation machine is like a super-precise, old-school calculator. It follows strict physics rules step-by-step to predict how sound travels. It's very accurate, but it's slow. It's like trying to predict the weather by calculating the movement of every single air molecule. It takes a long time, which makes creating the final image slow and expensive.

The New Idea: The "AI Predictor"

The authors of this paper asked: What if we taught a computer to guess the answer instead of calculating it from scratch every time?

They used a special type of Artificial Intelligence called a Fourier Neural Operator (FNO).

• The Analogy: Imagine you have a student who has watched thousands of weather simulations. They haven't memorized the physics formulas, but they have seen so many examples that they can look at the starting conditions and instantly say, "Oh, the wind will blow this way."
• How it works here: The AI was trained on thousands of examples of how sound waves travel. Once trained, it can predict how sound moves through the body almost instantly, without doing the heavy math calculations every time.

The Experiment: Teaching the AI to Solve the Puzzle

The researchers tested this new AI "Simulation Machine" in two ways:

1. Can it mimic the old machine?
   They asked the AI to predict sound waves and compared it to the slow, traditional calculator.

   • Result: The AI was incredibly accurate. It made very few mistakes, looking almost identical to the slow calculator's results. It was also much faster (about 8 times faster in their tests).

2. Can it help solve the "Inverse Problem" (the picture reconstruction)?
   This is the hard part. They used the AI to help rebuild the image of the hidden object.

   • The Challenge: Usually, to fix a blurry image, you need to know exactly how the math changes if you tweak the picture slightly (this is called a "gradient"). With the old calculator, figuring out these changes is slow.
   • The Solution: Because the AI is built with modern software (PyTorch), it has a superpower called Automatic Differentiation. It's like the AI has a built-in "undo" button that instantly tells you how the sound would change if you moved a pixel in the image. This makes the image reconstruction process incredibly fast.

The Results: Fast and Accurate

They tested this on different scenarios:

• Full View: Sensors all around the object (easy).
• Limited View: Sensors only on one or two sides (hard, like trying to see a face from only the side).

The findings were great:

• Accuracy: The pictures made with the AI were just as good as the pictures made with the slow, traditional method. Even when the sensors were in a "limited view" position, the AI didn't get confused.
• Speed: The AI calculated the necessary steps for the image reconstruction in a fraction of a second, whereas the traditional method took much longer.
• Generalization: They tested the AI on a shape it had never seen before (a "Shepp-Logan" phantom, which is a standard test pattern). The AI still worked well, proving it actually learned the physics, not just memorized the training pictures.

The Takeaway

This paper shows that we can replace the slow, heavy-duty physics calculators with a fast, trained AI to solve medical imaging problems.

• Old Way: "Let me calculate every single step of the sound wave..." (Slow, accurate).
• New Way: "I've seen this before; I know how the sound behaves. Let me guess the path instantly." (Fast, accurate).

This means that in the future, doctors might be able to get high-quality 3D images of the inside of a body much faster, potentially allowing for real-time imaging during surgeries or faster diagnoses. The AI acts as a "shortcut" through the math, saving time without losing quality.

Technical Summary

1. Problem Statement

Photoacoustic Tomography (PAT) is an imaging modality that reconstructs the initial pressure distribution ($p_0$) inside a target from ultrasound waves measured on the boundary. This is an inverse problem governed by the acoustic wave equation.

• The Challenge: Solving the inverse problem typically requires a "forward operator" ($K$) that simulates how the initial pressure propagates as ultrasound waves. Traditional methods use numerical approximations like the pseudospectral k-space method (implemented in tools like k-Wave). While accurate, these methods are computationally expensive, creating a bottleneck for iterative reconstruction algorithms that require thousands of forward and gradient evaluations.
• The Goal: The authors aim to replace the traditional numerical forward operator with a learned forward operator based on deep learning to accelerate the reconstruction process while maintaining accuracy.

2. Methodology

A. The Learned Forward Operator (Fourier Neural Operator)

Instead of solving the Partial Differential Equation (PDE) numerically at every step, the authors employ a Fourier Neural Operator (FNO).

• Architecture: A 3D deep learning network (2 spatial dimensions + 1 temporal dimension) designed to learn the solution operator of the acoustic wave equation.
• Input/Output: The network takes the initial pressure distribution ($p_0$) as input and outputs the full spatio-temporal ultrasound wave-field ($p_t$).
• Training: The FNO is trained on pairs of initial pressures and corresponding wave-fields generated by the high-fidelity pseudospectral k-space method.
• Hyperparameters: The network uses 2 Fourier layers, 4 channels, and 64 Fourier modes. It is trained using the ADAM optimizer.

B. Solving the Inverse Problem

The reconstruction is framed as a Bayesian inverse problem, seeking the Maximum A Posteriori (MAP) estimate.

• Objective Function: The authors minimize a cost function comprising a data fidelity term (difference between measured and predicted data) and a regularization term (Gaussian Ornstein-Uhlenbeck prior).
• Optimization Algorithm: The Broyden-Fletcher-Goldfarb-Shanno (BFGS) quasi-Newton method is used for minimization.
• Gradient Computation: A key innovation is the use of Automatic Differentiation (AD) (via PyTorch). Instead of deriving analytical gradients or using adjoint state methods for the k-space solver, the gradient of the objective function with respect to $p_0$ is computed directly through the FNO using torch.autograd.grad. This allows for efficient Jacobian-vector product calculations.

C. Experimental Setup

• Domain: 2D square domain ($10 \text{ mm} \times 10 \text{ mm}$).
• Phantoms:
  • Training/Testing: Blood vessel-like structures (from High-Resolution Fundus Image Database).
  • Generalization Test: Shepp-Logan phantom (out-of-distribution).
• Sensor Geometries: Full-view, two-side limited-view, and one-side limited-view.
• Noise: 1% Gaussian noise added to simulated data.

3. Key Contributions

1. Integration of FNO with Iterative Reconstruction: The paper demonstrates a framework where a learned neural operator replaces the traditional PDE solver within a gradient-based iterative optimization loop (BFGS).
2. Automatic Differentiation for PAT: It showcases the utility of automatic differentiation to compute gradients of the learned forward operator, eliminating the need for complex adjoint derivations for the neural network.
3. Generalization Capability: The study proves that the FNO, trained on vessel-like structures, can generalize to reconstruct Shepp-Logan phantoms (a different class of objects) without retraining.
4. Computational Efficiency: The approach significantly reduces the time required for forward simulations and gradient calculations compared to the state-of-the-art k-space method.

4. Results

Forward Simulation Accuracy

• The FNO approximated the photoacoustic wave-field with high accuracy.
• Relative Difference (RD):
  • Vessel phantoms (in-distribution): $6.16 \pm 2.11\%$.
  • Shepp-Logan phantom (out-of-distribution): $20.6\%$.
• Speed: Forward simulation took 0.057 seconds with FNO vs. 0.44 seconds with the k-space method (approx. 7.7x faster).

Inverse Problem Reconstruction

• Visual Quality: The MAP estimates obtained using the FNO were visually indistinguishable from the reference k-space method across all sensor geometries (full-view and limited-view). Both methods exhibited similar blurring in limited-view scenarios, which is a physical limitation of the data, not the algorithm.
• Quantitative Error (Relative Error - RE):
  • Vessel Dataset: FNO and Reference methods yielded nearly identical errors (e.g., Full-view: $4.63\%$ vs. $4.58\%$).
  • Shepp-Logan Dataset: Errors remained comparable (e.g., Full-view: $16.5\%$ vs. $15.5\%$).
• Convergence: The BFGS algorithm converged in approximately the same number of iterations for both methods, though the FNO sometimes reached a slightly higher final objective function value due to minor modeling errors.

Computational Performance

• Gradient Computation: This is the most significant gain.
  • Reference Method (k-Wave): ~0.47 seconds per gradient calculation (constant regardless of sensor count).
  • FNO Method: ~0.006 seconds per gradient calculation (approx. 70x faster).
  • The FNO gradient time remained low even as sensor geometry changed, whereas the k-space method remained slow.

5. Significance and Conclusion

The paper establishes that Fourier Neural Operators are a viable, computationally efficient alternative to traditional numerical solvers for the forward model in Photoacoustic Tomography.

• Efficiency: By leveraging automatic differentiation and learned operators, the computational cost of solving the inverse problem is drastically reduced, making iterative reconstruction feasible for larger datasets or real-time applications.
• Accuracy: The method maintains reconstruction quality comparable to the gold-standard k-space method, even for out-of-distribution targets and limited-view geometries.
• Future Outlook: The authors note that while the current study is 2D, the methodology is flexible. Future work involves extending to 3D (addressing memory constraints) and validating with experimental data to account for real-world factors like sensor directivity and frequency response.

In summary, this work bridges the gap between deep learning and classical inverse problems, offering a path toward faster, high-fidelity photoacoustic imaging.