A few months ago, I saw this excellent video about mechanistic interpretability of large language models. The video goes over a language model trained on an exceedingly simple task and inspects its weights and activations in detail to understand exactly what is happening under the hood. This appealed to me for two reasons:
1. I dislike that machine learning models are often black boxes, and I don’t love relying on them when I can’t understand how they work
2. The training data for this exercise was synthetically produced, so no need to wrangle some huge dataset
This got me thinking about what a similar approach could be in the computational geometry space. In the end, I settled on a simple neural network that accepts an array of samples of a 2D signed distance field (SDF) of a circle and infers the radius and center position of that circle. The training data is easy to generate synthetically, the network is so small that the whole thing could be trained on a very low-power computer, and the weights and activations should be simple to inspect.
But What is an SDF?
An SDF is a function that takes a point in space as input and spits out the minimum distance from that point to the surface of whatever geometry you’re interested in. The “signed” part comes from the fact that points outside the geometry are positive and points inside are negative. There are many primitives that have closed-form SDF evaluations, even some pretty complicated shapes have closed-form SDFs, but for the general case like meshes you’re stuck with a lot of computation to find this value. Here I’m using an SDF for a circle, which is one of the simplest ones around.
Building the Network
The network I’m using here consists only of a convolutional layer and a fully connected layer. The input is a 9×9 grid of sampled points centered at the origin, and the output is three numbers: the radius, the center x coordinate, and the center y coordinate. There is a ReLU activation after the convolutional layer. That’s it. That’s the whole network.
Now we need to actually get meaningful weights for this network. It’s all so small that it converges within about twenty training iterations, and the final result is able to infer the parameters of the shape with surprising accuracy. Let’s look at how the convolutional kernels evolve over training.
The majority of the value of the convolutional weights is concentrated in the corners, and these values seem to take shape pretty quickly after training starts. Most of the process is dedicated to decreasing the values of the irrelevant weights so they don’t get in the way of the components that are doing the heavy lifting. This makes sense: at the start of training, the network will take large steps towards its solution, then as time goes on it will fine-tune and narrow in on a precise solution. Interestingly, we also see some symmetry here: the first and third kernels are almost mirror images of each other.
Let’s move on to the fully connected layer. In practice this is a 27×3 matrix, but I’ve rearranged things so they match up with the previous layer for our understanding.
Once again, we see the weights taking shape within the first few iterations of training, then fine-tuning later on. And once again, there is a lot of symmetry for these weights. The first column corresponds to the radius output, and we see that the top square of weights and the bottom square are mirrors of each other. They also are negations of the corresponding convolutional kernels from the previous layer. The second and third columns form the x and y coordinates of the center of the circle, and we also see symmetry here: All three blocks are mirror images of each other, up to a sign. This would make sense, as the x and y coordinates work in essentially the same way, just along a different axis.
Peeking Activations
These weights only really tell half the story. How do these weights somehow magically spit out the correct parameters at the end of the network? It’s not like one of these layers is just a set of calipers that measures things about the circle it’s supposed to see. To understand that, we need to look at the activations. For this, I’ve created a test sequence that smoothly varies the radius, then the x coordinate, then the y coordinate of the input. While this changes, we can inspect the resulting activations.
Now we’re getting somewhere. The first convolutional channel generally measures how far the values are along the axis moving from bottom left to upper right. The second convolutional channel generally shows how negative or positive the function is. And the third channel behaves very similarly to the first, but from lower right to upper left.
The vast majority of the fully connected activations are concentrated in the center chunk of the fully connected weights. As the circle grows, the points on average become more negative, so it makes sense that the center chunk of the radius weights are looking at the negative sum of the values found at the four corners. The center chunks of the horizontal and vertical activations show a way to measure how far in a given direction the values are shifted: shifting in one axis leads to activations that cancel each other out, and a shift in the other axis leads to a corresponding shift in the activation values.
That reveals what the center chunks of fully connected weights are calculating, but that still leaves the question of what all of the other weights are working on. To evaluate this, let’s give each of the three outputs their own time-series plot and plot the contributions of each of the three “channels” separately.
The input here is sinusoidal, and the behavior we see appears to be a series of constructive and destructive waves. For the radius calculation, the second channel can pretty much exactly calculate the radius, up to a static offset. The bad news is that this channel generates a lot of motion later on in the plot where the radius should be constant. To fix this, the first and second channels combine to destructively interfere and cancel out the unwelcome shift. For the horizontal and vertical center calculations, we see a similar story: the second channel provides the majority of the energy, with the first and third channels stepping in to interfere and iron out the rest of the error.
Wrapping Up
Looking under the hood of how neural networks learn to process signed distance functions, we can see that they essentially learn to measure different aspects of the input, then combine these measures into meaningful outputs. In only a few training iterations, the model converges to a highly accurate solution that captures the essence of the problem.
Given the similarities between the weights learned here and a set of basis vectors, I’d be interested to see if initializing a set of weights with a set of orthogonal bases is a reasonable thing to do for these neural networks. For something like this, you could take a sampling of potential inputs and compute the eigenvectors of the distribution, or something similar.