JAX-BEM: Gradient-Based Acoustic Shape Optimisation via a Differentiable Boundary Element Method

This paper introduces JAX-BEM, a differentiable Boundary Element Method solver built on the JAX framework that matches the accuracy of traditional codes while enabling efficient gradient-based acoustic shape optimization and inverse problem solving.

The Gist

Imagine you are trying to design the perfect loudspeaker horn. You want it to spread sound evenly in a room, like a gentle rain, rather than spraying it in a harsh beam like a fire hose. Traditionally, engineers would guess a shape, test it, see it's too loud in the corners, tweak the shape, test it again, and repeat this hundreds of times. It's slow, expensive, and often relies on luck.

This paper introduces JAX-BEM, a new "super-smart" way to design these shapes using a technique borrowed from Artificial Intelligence (AI).

Here is the breakdown using simple analogies:

1. The Old Way vs. The New Way

• The Old Way (Trial and Error): Imagine trying to find the best route through a foggy maze. You take a step, hit a wall, go back, try a different direction, hit another wall, and so on. You are "blind" to the slope of the terrain.
• The New Way (Gradient-Based): Now, imagine you have a magical map that not only shows the maze but also gives you a compass that always points "downhill" toward the exit. This is what Gradient-Based Optimization does. Instead of guessing, the computer calculates exactly which tiny tweak to the shape will improve the sound the most.

2. The "Boundary" Trick (BEM)

To simulate sound, you usually have to fill the entire room with a 3D grid of tiny cubes (like a voxel game) to calculate how sound waves bounce around. This is heavy and slow.

The Boundary Element Method (BEM) used in this paper is like painting a room. Instead of filling the whole room with paint (meshing the air), you only paint the walls (the boundaries).

• Analogy: If you want to know how wind flows around a car, you don't need to simulate every drop of air in the sky. You just need to simulate the air touching the car's surface. BEM does this for sound, making the math much faster.

3. The "Magic Sauce" (JAX & Automatic Differentiation)

The real breakthrough here is using a tool called JAX. JAX is a software framework originally built to train massive AI brains (neural networks) that have billions of connections.

• The Problem: Usually, if you write a complex physics simulation (like sound waves), the computer can't easily figure out "If I move this wall 1 millimeter to the left, how does the sound change?" It's like trying to reverse-engineer a cake recipe by eating the cake.
• The Solution: JAX has a feature called Automatic Differentiation. It's like a "undo" button that remembers every single math step the computer took. If you ask, "How do I change the input to get a better output?", JAX instantly traces the steps backward to tell you exactly how to tweak the shape.

4. What They Actually Did

The researchers built a sound simulator that is fully differentiable. This means they could feed it a goal (e.g., "I want the sound to be loud in the center but quiet in the corners") and let the computer automatically reshape the loudspeaker horn to achieve it.

• The Result: They started with a simple cone-shaped horn. The computer, using this "magic compass," slowly warped the metal into a complex, curvy, organic-looking shape.
• The Outcome: The final shape was much better at spreading sound evenly than the original cone. It reduced "diffraction" (those annoying sound spikes that happen at sharp edges) and created a smoother sound field.

5. Why This Matters

• Speed: Because they used JAX, they could run these simulations on powerful graphics cards (GPUs), making the process 3 to 4 times faster than previous methods.
• Complexity: It allows engineers to optimize shapes that are too weird or complex for a human to design by hand.
• Future: While they tested this on loudspeakers, this same "magic" could be used to design better antennas for phones, quieter wind turbines, or even medical ultrasound devices.

Summary

Think of JAX-BEM as giving an engineer a self-driving car for sound design. Instead of manually steering the car (tweaking the shape) and hoping to get to the destination (perfect sound), you just tell the car where you want to go, and the car's AI (the gradient tracker) steers itself there, finding the smoothest, most efficient path through the landscape of physics.

Technical Summary

1. Problem Statement

Traditional numerical optimization for engineering structures, particularly in acoustics, faces significant challenges when dealing with multiple objectives and high-dimensional parameter spaces.

• Limitations of Current Methods: Domain-based methods (FEM, FDTD) require meshing the entire air volume, leading to a massive number of degrees of freedom that scale quadratically (2D) or cubically (3D) with domain size. While the Boundary Element Method (BEM) is more efficient (meshing only boundaries), traditional BEM solvers are not inherently differentiable.
• The Gap: Gradient-based optimization requires the ability to compute derivatives of the solution with respect to design parameters (e.g., geometry). Standard iterative solvers (like GMRES) used in BEM are difficult to differentiate because they involve loops and indeterminate iteration counts, making them incompatible with standard automatic differentiation (autodiff) frameworks.
• Goal: To develop a differentiable BEM solver that enables efficient, gradient-based shape optimization for acoustic devices (specifically loudspeaker horns) in unbounded domains.

2. Methodology

The authors developed JAX-BEM, a differentiable BEM solver built using the JAX automatic differentiation framework.

Core Technical Approach

• Framework: The solver refactors the open-source bempp operators into Just-In-Time (JIT) compiled, vectorized JAX functions. This leverages JAX's ability to track gradients through complex numerical algorithms.
• Handling Non-Differentiable Components:
  • Mesh Adjacency: Calculating mesh adjacency and selecting integration types (for singular/near-singular interactions) involves control flow (if statements) which is non-differentiable. The authors moved these calculations outside the optimization loop, computing them once per iteration before the differentiable sequence begins.
  • Implicit Differentiation: The BEM solver uses GMRES (an iterative linear solver). Unrolling GMRES for backpropagation is memory-intensive and inefficient. Instead, the authors used Implicit Function Theorem via JAX's custom_vjp (Vector-Jacobian Product).
    • Forward Pass: Solves $Ax = b$ using GMRES to get pressure $p_s$.
    • Backward Pass: Instead of unrolling the solver, it solves the adjoint system $A^H \lambda = g$ (where $g$ is the gradient of the loss). Gradients with respect to the matrix $A$ and vector $b$ are derived from $\lambda$, and then propagated to the mesh vertices via the chain rule. This treats the solver as a "black box" for autodiff.
• Hardware Acceleration: The method utilizes OpenXLA backends to run on CPUs, GPUs, and TPUs. The boundary-to-domain propagation phase is "embarrassingly parallel," making it highly suitable for GPU acceleration.

Optimization Setup

• Objective: Optimize the directivity of a loudspeaker horn.
• Parameters: Control points of splines defining the axial horn profile.
• Loss Function: Mean Squared Error (MSE) between the simulated sound pressure level (SPL) and target dB levels.
  • In-coverage region: Target $T_{in} = -3$ dB.
  • Out-of-coverage region: Target $T_{out} = -10$ dB.
• Optimizer: L-BFGS (via scipy.optimize.minimize).

3. Key Contributions

1. First Differentiable BEM Implementation: Demonstrates a fully differentiable Boundary Element Method solver capable of backpropagating gradients from the domain solution all the way back to the mesh vertices.
2. Implicit Differentiation for Iterative Solvers: Successfully applies implicit differentiation to the GMRES solver, bypassing the need to unroll iterative loops, which drastically reduces memory overhead and enables the use of standard iterative solvers in gradient-based optimization.
3. Hardware Acceleration: Shows that BEM, traditionally CPU-bound due to dense matrix operations, can be significantly accelerated on GPUs (3-4x faster than CPU-based bempp on standard hardware).
4. Application to Acoustic Design: Successfully applies the method to a complex, real-world problem: optimizing a loudspeaker horn for constant directivity across a wide bandwidth (4 kHz – 18 kHz).

4. Results

• Validation (Accuracy):
  • Tested against the analytic Mie series solution for scattering from a rigid sphere.
  • Error: JAX-BEM matched the error of the established bempp code. Mean Absolute Error (MAE) relative to analytic solutions ranged from $4 \times 10^{-2}$ (coarse mesh) to $4 \times 10^{-4}$ (fine mesh).
  • Precision: Despite JAX defaulting to 32-bit precision (complex64) vs. bempp's 64-bit (complex128), there was no significant difference in error, indicating that discretization error dominates over floating-point precision issues.
• Performance (Speed):
  • JAX-BEM on CPU was 3–4x faster than bempp (which uses Numba).
  • GPU Speedup: On a Radeon GPU, the solver became faster than the CPU version for meshes with >4,000 elements.
• Optimization Outcome:
  • Initial State: A simple conical horn showed significant diffraction above 8 kHz.
  • Optimized State: The optimizer generated a complex mouth shape with a smooth termination.
  • Directivity: Diffraction was reduced, and horizontal directivity became more constant. However, widening the vertical directivity to the target increased out-of-coverage SPL at certain frequencies, highlighting the trade-offs in multi-objective optimization.
  • Memory: A single iteration (20 frequencies) consumed ~100 GB of RAM in 32-bit precision, highlighting the $O(N^2)$ scaling challenge of BEM.

5. Significance and Future Work

• Significance: This work bridges the gap between machine learning (automatic differentiation) and classical computational physics (BEM). It enables the use of powerful gradient-based optimizers for inverse problems and geometry design in acoustics, electromagnetics, and potentially other wave physics.
• Challenges Identified:
  • Memory Usage: The $O(N^2)$ scaling of BEM combined with the memory overhead of storing the computational graph for autodiff is a bottleneck. The authors suggest H-Matrix compression as a potential future solution if implemented in a differentiable framework.
  • Solver Efficiency: Re-solving the entire system from scratch for every geometry update is inefficient. Future work should explore "warm-starting" GMRES with previous solutions.
• Future Directions:
  • Validation with physical 3D-printed devices.
  • Comparison against non-gradient-based optimization methods.
  • Extension to electromagnetic wave simulations.

In conclusion, JAX-BEM represents a paradigm shift in acoustic simulation, moving from "simulation for analysis" to "differentiable simulation for design," enabling faster and more complex optimization of acoustic geometries.