Building an MLP with flax.linen.Dense

Welcome back to the second lesson of Beyond Pure JAX: Flax & Optax for Elegant ML! You've made excellent progress so far. Having mastered the foundational concepts of Flax modules with setup() and __call__() methods in our previous lesson, you're now ready to explore one of Flax's most essential building blocks: the Dense layer.

While creating custom modules gives us a deep understanding and flexibility, real-world machine learning relies heavily on proven, optimized components. The flax.linen.Dense layer is precisely such a component — a fully connected layer that handles weight initialization, bias terms, and efficient linear transformations automatically. Today, we'll construct a complete Multi-Layer Perceptron (MLP) using these built-in Dense layers, learning how to compose them elegantly while exploring two distinct approaches to module definition.

By the end of this lesson, you'll confidently build neural networks using Flax's high-level abstractions, understanding both the setup() method and the powerful @nn.compact decorator. We'll apply these techniques to construct an MLP capable of learning the XOR function, demonstrating how Flax simplifies complex model construction while maintaining JAX's functional programming principles.

The flax.linen.Dense layer represents a fundamental building block in neural networks, implementing the ubiquitous linear transformation:

where $W$ is the weight matrix, $x$ is the input vector, and $b$ is the bias term. Unlike our custom implementations from the previous lesson, Dense layers come with sophisticated initialization strategies, automatic shape inference, and optimized computation paths.

What makes Flax's Dense layer particularly elegant is its simplicity of use combined with powerful underlying functionality. When we specify nn.Dense(features=3), we're creating a layer that will automatically determine input dimensions from the first forward pass, initialize weights using proven strategies like Xavier/Glorot initialization, and handle the linear transformation efficiently using JAX's optimized operations.

The features parameter determines the output dimensionality, that is, how many neurons this layer contains. For our XOR problem, we might use Dense(3) for a hidden layer with three neurons and Dense(1) for the output layer producing a single prediction. Flax handles all the complexity of weight matrices, bias vectors, and shape compatibility automatically, allowing us to focus on architecture design rather than implementation details.