Concise Implementation of Linear Regression

Dive into Deep Learning · §2.5

The same model, the concise way
batteries-included layers, losses, and optimizers replace the hand-rolled parts.

From hand-rolled to high-level

Motivation

Last section we wrote every piece by hand: the weight vector, the forward pass, the squared error, the update step.

Those pieces are so universal that frameworks ship them, tuned and tested. We swap each one for its built-in counterpart:

Layer replaces w, b · loss replaces our squared error · optimizer replaces the update loop.

By hand	Built-in
w, b	a layer
MSE math	a loss
update step	an optimizer

01

The Model

a single linear layer

The layer already is the model

The Model

What we hand-rolled as w, b, and a matrix–vector product, every framework ships as a fully connected layer: each input wired to the one output: exactly the picture of linear regression.

The layer owns its parameters. We no longer allocate them, initialize them, or even know their shapes ahead of time.

One fully connected layer with a single output is linear regression.

One layer, with explicit sizes and randomness

The Model

nnx.Linear(num_inputs, 1, rngs=rngs) is the whole model. NNX has no lazy mode, so we state the input width and hand the layer an explicit RNG stream — and its parameters exist, on the module, as soon as the constructor returns:

class LinearRegression(d2l.Module):
    """The linear regression model implemented with high-level APIs."""
    def __init__(self, num_inputs, lr, rngs=None):
        super().__init__()
        self.save_hyperparameters(ignore=['rngs'])
        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
        self.net = nnx.Linear(
            num_inputs, 1, kernel_init=nnx.initializers.normal(0.01),
            rngs=rngs)

Two JAX signatures — explicit shapes, explicit randomness. The weights then live on the module, just as in the other frameworks.

The forward pass is a one-liner

The Model

forward just calls the layer. All the matrix–vector arithmetic we wrote by hand now lives inside it:

@d2l.add_to_class(LinearRegression)
def forward(self, X):
    return self.net(X)

02

Loss & Optimizer

two more pieces, off the shelf

Loss: built-in mean squared error

Loss & Optimizer

The framework’s MSE replaces our hand-written squared error:

@d2l.add_to_class(LinearRegression)
def loss(self, y_hat, y):
    return d2l.reduce_mean(jnp.square(y_hat - y))

It omits the \tfrac{1}{2} factor we used by hand, and averages over the minibatch by default.

Optimizer: minibatch SGD in one call

Loss & Optimizer

The update loop becomes a single optimizer object, handed the parameters and the learning rate:

@d2l.add_to_class(LinearRegression)
def configure_optimizers(self):
    return optax.sgd(self.lr)

The same optim/Trainer family also gives momentum, Adam, and more by swapping one line.

03

Training

the scaffold never changed

The same Trainer drives it all

Training

Our Trainer, Module, and DataModule from the object-oriented-design section don’t care that the model is now a built-in layer.

The training loop is identical to the from-scratch version.

Fit: same data, same curve, a fraction of the code

Training

Same synthetic data, same ten epochs, same fit call as the linear-regression-from-scratch section:

model = LinearRegression(2, lr=0.03)
data = d2l.SyntheticRegressionData(w=d2l.tensor([2, -3.4]), b=4.2)
trainer = d2l.Trainer(max_epochs=10)
trainer.fit(model, data)

Nothing about the training run can tell the two implementations apart: only the amount of code we wrote changed.

Where the parameters live now

Training · payoff

They no longer hang off our class as self.w, self.b; they live inside the layer, so get_w_b reaches through net:

@d2l.add_to_class(LinearRegression)
def get_w_b(self):
    return self.net.kernel[...], self.net.bias[...]

w, b = model.get_w_b()

print(f'error in estimating w: {data.w - d2l.reshape(w, data.w.shape)}')
print(f'error in estimating b: {data.b - b}')

error in estimating w: [ 8.2969666e-05 -2.4771690e-04]
error in estimating b: [0.00036526]

Same verdict as the linear-regression-from-scratch section: the true \mathbf{w}^* = [2,-3.4], b^* = 4.2 recovered to a few 10^{-4}. The built-in pieces really do compute the same thing our hand-rolled ones did.

Summary

Wrap-up

• From scratch showed what happens; concise is what we actually use day to day.
• A single layer stands in for w, b; a built-in loss and optimizer replace the rest.

• The Module / Trainer / DataModule scaffold is unchanged; only the model’s internals got shorter.
• Same minibatch loop, same convergence: ~5 lines of model code, error order 10^{-4}.