Neural Texture Compression in Metal 4

Introduction

This is an implementation of Neural Texture Compression (NTC) which is based on the NVIDIA research paper: Random-Access Neural Compression of Material Textures. I was at the point where I wanted to do texture compression for my custom engine, and after skimming through this research, I was impressed with the compression-quality results. So I went down the rabbit hole of machine learning.

This post will serve as extended documentation of MetalNTC and I will try not to re-explain stuff that can be found in the original research, but rather cover details that haven’t been mentioned, to help give you an overview of the whole pipeline of this project. I will discuss positives/negatives and fail cases, as well as use cases of the NTC, since it for now can’t entirely replace traditional texture sampling.

Real time inference is costly, but in a later section I will discuss how NTC can provide value even on an M1 Air.

Creating the Ground Truth

The Multilayer Perceptron (MLP) will be able to decompress any texture and mip level of the material texture set. Hence it needs the mips for each one of the source textures. Conveniently, my last post was on a Single Pass Downsampler which I could reuse. The only difference now is that it receives a texture2d_array, and every thread loops over the number of materials the texture set includes.

Training

Forward Pass

1. The neural network is using latent grids with learned features. Each mip level is using a pair of latent grids, where this pair can be used by many mip levels. This sharing of features lowers the storage cost of a traditional mipmap chain. The grid G0 is at higher resolution, which helps preserve high frequency details, while G1 is at a lower resolution, helping with the reconstruction of low frequency content. Both of them are bilinearly sampled and their features are later concatenated and not blended. The choice of which grid pair that a mip should use is empirical.
2. During training uniform noise is applied that simulates what would happen if the parameters were quantized. Quantization produces an error that is approximately uniformly distributed, hence the [-q/2, q/2) noise is applied. The actual quantization takes place at a later stage, but the neural network must be trained on the expected precision loss. The round() in the quantization would create zero gradients during training, since a small nudge does nothing at all. The gradient is multiplied by the chain rule during backpropagation, so the latents would receive nothing and never get trained.
3. Concatenation is the process that will construct the input layer of the MLP. As mentioned before, each latent grid’s features are separate from each other and each one has 16. The positional encoding (PE) gives the network a signal that the grid’s resolution cannot physically encode. For the case of a 2048x2048 texture, the G1 grid at 256x256 (which is the biggest upsampling) is 8 times smaller than the source. One texel in G1 spreads across an 8x8 block of output texels, so within that block the G1 features vary only as a linear blend of the same four stored texels. PE is what varies inside the block. Each latent grid pair is used for at least 3 texture mips. All 44 other inputs are identical across the mips that latent pair serves, so the network would be handed the same input and asked for three or more different images. The normalized LOD is the only thing that is different. The padding is needed because the input layer needs to be a multiple of 16 for the tensor ops.
4. The RTXNTC which is the repository of the research, is currently using 64-48-32 hidden layers. However I went with their original 64-64 hidden layers as it was simpler to implement and I’m leaving the new design for the future. The hardGELU is a fast piecewise quadratic approximation of the standard GELU activation.
5. Calculate loss against the ground truth. Channels of different textures have different importance. For example the albedo’s error gets a bigger say in how much it drives the update, whereas textures like the displacement map get smaller weight. Because of the fixed number of features per grid, the weights need to favor the textures whose error shows up in the final image.

Backpropagation

After loss, backpropagation runs backwards on the same layers and same weights but transposed. Forward reads 4 corners and produces 1 feature (gather). Backward takes 1 gradient and writes 4 corners (scatter). The size of the batch is 4096. Each thread in the batch is accumulating its gradients in the shared gradient buffer. Metal doesn’t have float atomics, so integer atomics are used with a scale factor. The layer weights and biases are accumulated by every thread of the batch, but each thread accumulates only the latent features that it touches. Instead of plain stochastic gradient descent (SGD), the parameters are updated using the Adam stochastic optimizer. The latent grid gradients are very sparse, so there can be cases where a specific texel won’t receive a gradient for many steps and eventually get one contribution from a single thread. Adam stores momentum for every parameter, so an untouched parameter will still keep descending for a number of steps. The squared gradient average helps normalize the learning rate between different parameters. For example MLP weights are hit by all 4096 threads every step, so their accumulated gradients are large unlike latent grid features.

Quantization

Actual quantization happens right after the training (with the simulated noise) has finished. The grids are quantized down to 4 bits per feature, therefore there are 16 representable values.

Τhe example below shows the error for the training value 0.1372 as in Figure 2.

Quantization example

Rounding is where precision is lost, but it is necessary for the value to fit in 4bits. Clamping keeps the value inside the bounds (clamping also takes place during the Adam step as in the table below). Decode is using the offset to convert the value from unsigned back to signed and after scaling by q, the error is “decoded - trained value”.

Following the paper, the training runs again for 5% of the original steps, but this time with the decoded latent features frozen. Only the MLP weights are updated to retrain on the rounded latent features.

Inference

Latents as Textures

The latent grid features are stored, among others, in the .ntc binary file after the training phase. Since the latent grids will be bilinearly sampled during the inference, the hardware accelerated linear interpolation can be taken advantage of, if the features are stored in textures. However, each latent grid texel is storing 16 features and the rgba4Unorm format (in Metal there is only abgr4Unorm but behaves like rgba4Unorm) can store 4 features. That can be solved by giving each texture 4 slices as seen in Figure 3.

Decoding Pixels

The latent grid features are loaded from the .ntc as a texture2d_array and the rest of the MLP weights as a buffer. The rest of the input parameters in Figure 4 are computed and concatenated during runtime. During decoding the forward pass will decode all of the channels in the material texture set, therefore for cases where only a texture needs to be inferred, like in a depth prepass, it is better to use traditional sampling there. The output layer is 16 channels, which should be enough for most of the material texture sets.

Stochastic Trilinear Interpolation

On current hardware, inference cost is very high even with tensor ops, and trilinear interpolation to blend mips requires inferring twice. The workaround is to stochastically choose between the mips and accumulate the results across frames.

The selection of mip level is calculated in the fragment shader using the dfdx(uv) and dfdy(uv) derivative functions.

The variable base is used for the base mip level to be inferred and the fractional part is the probability of sampling the base or base + 1 mip level. Jitter varies per pixel using blue noise and per frame using the golden ratio offset. In this example in the table above, mip level 4 has a 70% chance of being inferred on this frame.

Note

In MetalNTC-Renderer a very simple temporal accumulation filter is used just to resolve the flickering, but in a real application a good TAA solution or similar would be needed.

Source code

• MetalNTC
• MetalNTC-Renderer

References

• Random-Access Neural Compression of Material Textures
• Metal Shading Language Specification