Simple Linear Model

A linear model is the smallest trainable network that is still useful: one matrix multiply, one bias vector, and a softmax that turns scores into probabilities. Every later lesson — CNNs, Inception, fine-tuning — reuses the same training loop with a bigger function in the middle. Start here to see the whole pipeline in about twenty lines.

The idea

Give each class a weighted sum of the input features. For an input vector $x \in \mathbb{R}^{D}$ and $K$ classes, the model computes

W has shape [D, K], b has shape [K], and the softmax is not a layer with weights — it is a fixed function. With no hidden activation, the decision boundary between any two classes is a hyperplane, so the model can only carve the input space into $K$ convex regions.

For MNIST, D = 28 × 28 = 784 and K = 10, so the model has 784 × 10 + 10 = 7850 parameters. That is all it needs, and all it can use.

Worked example

Take a two-feature input $x = [2, -1]$ and two classes with weights $w_3 = [0.5, 0.25]$, $w_5 = [-0.5, 0.75]$ and biases $b_3 = 0.1$, $b_5 = -0.1$.

• $\text{logit}_3 = 2(0.5) + (-1)(0.25) + 0.1 = 0.85$
• $\text{logit}_5 = 2(-0.5) + (-1)(0.75) - 0.1 = -1.85$

Softmax gives $e^{0.85} = 2.3396$ and $e^{-1.85} = 0.1572$, so $P(3) = 2.3396 / 2.4968 \approx 0.937$. The cross-entropy loss for the true label 3 is $-\ln(0.937) \approx 0.065$.

In code

import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(28, 28)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10),  # raw logits, no activation
])
model.compile(
    optimizer=tf.keras.optimizers.Adam(0.001),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
)

model.fit(x_train, y_train, epochs=5, batch_size=128, validation_split=0.1)

Keep the output as raw logits and let from_logits=True apply the softmax internally; that is numerically safer than a softmax followed by a log. The Flatten is a deliberate choice — a dense layer over MNIST pixels has no idea which pixels are neighbors.

Check yourself

1. Why does Dense(10, activation="softmax") paired with SparseCategoricalCrossentropy() (no from_logits) risk numerical problems?
2. How many parameters does the model have if you insert Dense(64, activation="relu") before the output layer?
3. What class of decision boundaries can this model represent, and what does that forbid on a dataset like MNIST?

Key takeaways

• A linear model is xW + b plus a softmax; the softmax itself has no parameters.
• MNIST linear shape: 784 → 10, which is 7850 parameters.
• from_logits=True is the numerically stable way to pair raw logits with cross-entropy.