Ray Tracing with Neural Networks

Introduction

Accurate simulation of imaging systems is crucial for enabling cost-effective and rapid design processes and training machine learning algorithms. Such simulation tools typically incorporate the full imaging pipeline, including the scene, camera optics, and image sensor. However, the proprietary nature of lens designs often leads manufacturers to withhold specific details, posing a challenge for accurate simulation. To address this challenge, Goossens et al. (2022) propose a solution that leverages Zemax black-box lens models. At the core of this method is the Ray Transfer Function (RTF), which characterizes how the position and angle of a ray entering a lens determine its exit position and angle, providing a practical approach to lens modeling without exposing proprietary design information. Goossens et al., modeled the RTF using a polynomial fitting and showed promising results as in [1]. The idea of this project is to explore the potential of neural networks to approximate the ray transfer function and enable efficient, high-fidelity ray tracing without the need for detailed lens specifications. As neural networks have demonstrated remarkable capabilities in learning complex nonlinear functions, making them a promising candidate for modeling the ray transfer function [2].

Background

Dataset Generation using Zemax

Polynomial Estimation of Ray-Transfer Functions

In a general form, the position of each input ray can be represented in six variables, $x,y,z,dx,dy,dz$. Similarly, the position of each output ray is represented as $x^{\prime },y^{\prime },z^{\prime },d{x}^{\prime },d{y}^{\prime },d{z}^{\prime }$. The polynomial-based approach aims to model this mapping between input and output ray positions using a set of polynomial equations:

$\hat{\beta }=\underset{{\beta }_0,{\beta }_1,{\beta }_2,\dots ,{\beta }_k}{\arg \min }\Vert [1,x,x^2,x^3,\dots ,x^k{]}^T[{\beta }_0,{\beta }_1,{\beta }_2,\dots ,{\beta }_k]-x^{\prime }{\Vert }_2^2$ where $x$ here is a column vector storing the "x" positions of all rays. The overall polynomial is multivariate but the logic is the same for $x,y,z,dx,dy,dz$. The user must specify the order of polynomial (denoted by $k$), as reported in [1], the fitting error varies as the polynomial degree changes (see the figure below).

Methods

Network Architecture

As a starting point, we are provided with the following GitHub repo: https://github.com/ISET/mlp-tiny, where a multi-layer perceptron (MLP) was implemented but the complete training/testing code was missing. We implemented the complete training/testing code and put all free parameters in a configuration file: https://github.com/ISET/mlp-tiny/blob/ray_tracing/mlp%20scripts/train_prediction.sh

The provided MLP has 16 fully connected layers (each is followed by batch normalization and ReLU activation function), as shown in the following figure.

However, one issue we found (see next section), was that the network struggled with predicting values in varying ranges, i.e., the coordinates can have significantly different maximums and minimums. To address the issue, we used positional encoding [3] that remapped the coordinates into [-1, 1] with a series of sinusoid functions at different frequencies, e.g., $x,sin(x),sin(2x),\dots ,sin(2^{p}x)$.

We integrated the positional embedding technique into the MLP and the resulting architecture is shown below:

where we found using separate branch to pre-process spatial and directional information and later merge them together to be the most effective. The modified architecture even has fewer trainable parameters than the MLP provided.

Loss Function and Training Details

We used binary cross-entropy loss (BCE, see https://pytorch.org/docs/stable/generated/torch.nn.BCELoss.html) for ray classification and mean squared error (MSE, see https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html) for ray tracing.

We set the hidden layer sizes in the MLP to be 128 units. The learning rate was set to 0.005 and warmed up over 40 epochs. A cosine decay schedule was employed to update the learning rate over training after the warm up. The Adam optimizer was used with a weight decay of 0.00001 [4]. We trained the MLP for 500 epochs with early stopping triggered if the training loss stopped decreasing. The dataset was partitioned into 60% training and 40% testing subsets.

Results

Ray Tracing

Visually, we found the training loss converged faster when using the "NeRF" style MLP [5] with positional encoding, compared to the provided "LargerMLP". Such observations were consistent with different wavelengths.

Our analysis revealed the following pattern for mean squared error on the test dataset: PolyFit model (d=2) has larger MSE than the larger MLP, which in turn surpassed the PolyFit model (d=5), followed by the NeRF approach, and finally the PolyFit model (d=8) again. Notably, a lower mean squared error indicates superior performance. In short: PolyFit (d=2) > LargerMLP > PolyFit (d=5) > NeRF > PolyFit (d=8).

References

[1] Thomas Goossens, Zheng Lyu, Jamyuen Ko, Gordon C. Wan, Joyce Farrell, and Brian Wandell, "Ray-transfer functions for camera simulation of 3D scenes with hidden lens design," Opt. Express 30, 24031-24047 (2022).

[2] Ian Goodfellow and Yoshua Bengio and Aaron Courville, "Deep Learning," MIT Press (2016), http://www.deeplearningbook.org.

[3] Position Embeddings. https://paperswithcode.com/methods/category/position-embeddings.

[4] Diederik P. Kingma, Jimmy Ba, "Adam: A Method for Stochastic Optimization," in ICLR, 2015.

[5] Ben Mildenhall and Pratul P. Srinivasan and Matthew Tancik and Jonathan T. Barron and Ravi Ramamoorthi and Ren Ng, "NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis," in ECCV, 2020.