Fast, Differentiable, GPU-Accelerated Ray Tracing for Multiple Diffraction and Reflection Paths

Technical Summary: Fast, Differentiable, GPU-Accelerated Ray Tracing for Multiple Diffraction and Reflection Paths

Problem Statement
Determining ray paths between a transmitter (TX) and receiver (RX) in complex environments typically involves exhaustive Ray Tracing (RT) algorithms that explore all combinations of object interactions. While Fermat's principle suggests that ray paths correspond to local extrema of path length (or total traversal time), existing optimization-based approaches face significant challenges in modern computing contexts:

1. Scalability on GPUs: Most current methods rely on extensive branching logic (e.g., separating reflection-only and diffraction-only paths) or variable problem dimensions based on interaction sequences. This prevents efficient uniform execution across parallel GPU threads.
2. Differentiability: Applications such as inverse design and network optimization require efficient gradient computation. Standard Automatic Differentiation (AD) applied to iterative solvers (like Newton methods) is computationally expensive and memory-intensive because it must backpropagate through every solver iteration.
3. Unified Handling: Existing solvers often treat reflections (via the Image Method) and diffractions (via minimization) as distinct processes, failing to support arbitrary sequences of mixed interactions efficiently.

Methodology
The authors propose a unified optimization framework that reformulates ray path tracing as a minimization of total path length, solvable via convex optimization techniques tailored for GPU acceleration.

• Unified Parametric Formulation: The problem is expressed in a parametric space where interaction points $x_i$ are affine transformations of variables $t_i$ ($x_i = A_i t_i + b_i$). By standardizing the dimensionality of these variables (using $d=2$ for both planar reflections and edge diffractions, where diffraction edges utilize a zeroed second basis vector), the method creates a uniform problem shape. This allows for vectorized computation across arbitrary sequences of reflections and diffractions without branching.
• Iterative Solver: The minimization of the path length objective function $L(T; A, B)$ is solved using the Broyden-Fletcher-Goldfarb-Shanno (BFGS) quasi-Newton algorithm. Unlike custom Newton methods used in prior work, BFGS avoids computing the Hessian matrix directly, which can be ill-conditioned near diffraction points.
• Fixed-Point Step Size: Instead of traditional inexact line searches, the authors employ a fixed-point iteration scheme to determine the optimal step size $\alpha$. This scheme converges rapidly (often within one iteration) and ensures uniform execution time across all threads in a batch, a critical requirement for GPU performance.
• Implicit Differentiation: To compute gradients with respect to problem parameters (e.g., object positions) without backpropagating through the entire solver history, the authors utilize the Implicit Function Theorem. By assuming the solver has converged to a stationary point where $\nabla_T L = 0$, they derive the total derivative $\partial T^*/\partial \theta$ by solving a linear system involving the Hessian of the objective function. This avoids the memory and computational overhead of differentiating through solver iterations.

Key Contributions

1. Generic Differentiable Solver: A single algorithmic procedure capable of handling arbitrary numbers and orders of reflections and diffractions, eliminating the need for separate solvers for different interaction types.
2. GPU-Accelerated Architecture: A formulation that maintains consistent problem dimensions, enabling efficient parallel execution on modern GPU architectures via libraries like JAX and DrJIT.
3. Efficient Gradient Computation: The implementation of implicit differentiation, which significantly outperforms standard reverse-mode AD in terms of speed and memory usage while providing exact gradients under convergence assumptions.
4. Open-Source Implementation: A clean implementation available in JAX, integrated into the DiffeRT RT tool.

Simulation Results
The method was evaluated on an NVIDIA GeForce RTX 3070 GPU, processing 1,000 parallel ray paths with up to 5 interactions ($n=1$ to $5$).

• Convergence and Accuracy: The proposed solver achieved convergence rates comparable to specialized Newton methods (Carluccio-Albani) and significantly outperformed standard Gradient Descent (which failed to converge in many cases). For mixed reflection-diffraction scenarios, the method was the only one capable of solving the problem, as other methods could not handle the variable dimensionality.
• Performance vs. Image Method: In pure reflection scenarios, the Image Method (IM) remained the fastest and most accurate. However, the proposed solver offered superior scalability and flexibility for mixed scenarios where the Image Method is inapplicable.
• Differentiation Efficiency: When computing gradients, the implicit differentiation approach was nearly 10 times faster than standard Automatic Differentiation. This performance gain remained constant regardless of the number of solver iterations or interaction complexity.

Significance and Claims
The paper claims that this work bridges the gap between physical ray tracing principles and modern differentiable programming. By unifying reflection and diffraction into a single convex optimization problem and leveraging implicit differentiation, the method enables new possibilities in inverse design and wireless propagation modeling that were previously hindered by the lack of scalable, differentiable solvers. The authors note that while the method is slower than the Image Method for pure reflection cases, its ability to handle arbitrary mixed sequences and its superior gradient computation efficiency make it a robust solution for complex, large-scale optimization problems. Future work aims to reformulate the problem as a Second-Order Cone Program (SOCP) to further improve convergence properties.