Stockfish NNUE Documentation

Original: nnue.md in official-stockfish/nnue-pytorch

Preface

What this document covers:

• Technical content
• A detailed description of NNUE and its principles
• A quick refresher on linear algebra
• Input definitions and factorization
• Components (layers) suitable for NNUE networks
• Inference code and optimization
• The mathematics of quantization and its implementation
• Nearly production-ready optimized code
• A PyTorch trainer implementation (+ important CUDA kernels)
• Architecture considerations and history

What this document does not cover:

• A network training tutorial (see the wiki)
• Datasets, optimizers, hyperparameters
• Logs of experimental results

Basics

What is NNUE?

NNUE (ƎUИИ Efficiently Updatable Neural Network) is, broadly speaking, a neural network architecture that takes advantage of the minimal changes in the network input between consecutive evaluations. It was invented by Yu Nasu for integration into the shogi program YaneuraOu, developed by Motohiro Isozaki in May 2018, and was later ported to the chess engine Stockfish by Hisayori Noda in June 2019. It is also applicable to many other board games and may even find use in other domains. NNUE follows these principles:

• The network should have a relatively small number of non-zero inputs.
• The inputs should change as little as possible between consecutive evaluations.
• The network should be simple enough to allow low-precision inference in the integer domain.

Following the first principle means that as the network grows, the input must become sparse. The best current architectures have an input sparsity of about 0.1%. A small number of non-zero inputs places a low upper bound on the time required for a full evaluation. This is the main reason NNUE networks can be very large and still be evaluated extremely quickly.

Following the second principle (provided the first is followed) creates a way to update the network (or at least its expensive part) efficiently, instead of re-evaluating the whole network. This exploits the fact that a single move changes the board state only slightly. It is less important than the first principle and entirely optional for an implementation, but implementations that do exploit it still see a measurable improvement.

Following the third principle enables maximum performance on commodity hardware and makes the model particularly well suited to low-latency CPU inference, which is required of a traditional chess engine.

Overall, the NNUE principles also apply to expensive deep networks, but they shine in fast, shallow networks suited to low-latency CPU inference without batching or accelerators. The target performance is millions of evaluations per second per thread. This is an extreme use case that demands extreme solutions, most importantly quantization.

Quantization 101 and why it matters

Quantization is the process of changing the domain of a neural network model from floating-point numbers to integers. NNUE networks are designed to be evaluated quickly in a low-precision integer domain and to make full use of the int8/int16 performance of modern CPUs. Floating point is not an option for achieving maximum engine strength, because it sacrifices too much speed for too little gain in accuracy (although some engines use floating-point representations for their simplicity). Quantization inevitably introduces error, which accumulates more the deeper the network is, but for the relatively shallow NNUE networks this error is negligible. Quantization is described in detail later in this document. Until then, the document uses floating-point numbers rather than integers; this will only matter when we actually get to code optimization. The purpose of inserting this section here is to make the reader aware of NNUE's end goal, since it is the biggest factor shaping NNUE models and determines what is and is not possible.

What layers are useful in NNUE?

NNUE relies on simple layers that can be implemented with simple arithmetic in a low-precision setting. This means that linear layers (fully connected, essentially matrix multiplication) and ClippedReLU (clamp(0, 1)) layers are particularly well suited. Pooling layers (multiplication/average/max) or approximations of more complex activation functions such as sigmoid are also usable, but are not commonly used.

Typically such networks are kept shallow (2-4 layers), because most of the knowledge is stored in the first layer (where input sparsity is exploited to keep performance up), and after the first layer the network needs to reduce its width sharply (the benefit of the deeper, later parts of the network would be dominated by the cost of the large earlier layers) to maintain performance requirements.

Linear layer

A linear (fully connected) layer is just a simple matrix multiplication. It can be implemented efficiently, supports sparse inputs, and provides good capacity. It takes in_features values as input and produces out_features values as output. The operation is y = Ax + b, where:

• x: the input column vector of size in_features
• A: the weight matrix of size (out_features, in_features)
• b: the bias column vector of size out_features
• y: the output column vector of size out_features

Linear layer with sparse inputs

The multiplication Ax can conceptually be simplified to "if x[i] is non-zero, take the i-th column of A, multiply it by x[i], and add it to the result". It should now be obvious that whenever an element of the input is zero, we can skip processing the corresponding whole column of the weight matrix. This means we only need to process the columns of A that correspond to the non-zero values in the input vector. Even though the weight matrix may have tens of thousands of columns, we only look at a handful of them per position! This is why the first layer can be so large.

Clipped ReLU layer

This is an activation function based on the ordinary ReLU, except that it is bounded at both the lower and the upper end. The formula is y = min(max(x, 0), 1).

The purpose of this layer is to add non-linearity to the network. If there were only linear layers, they could all be collapsed into one, since the matrices can simply be multiplied together.

Ideally ClippedReLU would be replaced by ReLU, but aggressive quantization requires reducing the dynamic range of the hidden-layer inputs, so clamping the values at 1 is important for performance.

Sigmoid

This is a smooth activation function, the counterpart of the [clipped] ReLU. The formula is y = 1/(1+e^-kx), where k is a parameter that determines how much the shape is "stretched".