Build a digit classifier head

Learn it

Words to know

logit

A raw score a model gives to one choice, before the scores are turned into probabilities.

argmax

The position of the largest value in a list.

feature

One of the numbers that describe an example, such as one pixel of an image or one word count.

dot product

Multiply two lists of numbers position by position and add up the results. It gives one number.

How it works

1. X has shape [examples, features] and W has shape [features, classes]. X @ W takes the dot product of each row of X with each column of W. That gives every example one score per class: shape [examples, classes].
2. Add bias, one number per class. PyTorch adds it to every row.
3. argmax(dim=1) gives, for each row, the position of its largest score. That position is the predicted class.

Worked example

Input
X = [[1, 2], [3, 0]]
W = [[1, 0, 2], [0, 1, 1]]
bias = [0, 0, -4]

Output
[1, 0]

Class 2 had the highest raw score for both examples. Its bias of -4 pulls it down, and a different class wins each time.

Why it matters

A GPT ends the same way: a linear layer that gives one score per token in the vocabulary. Picking the largest is the greedy decoding you write in chapter 7.

Your turn

Using PyTorch, compute class scores with one linear layer and return the predicted class of each example.

X

a table of shape [examples, features]

W

the weights, shape [features, classes]: one column per class

bias

one number per class

Returns

A Python list with the index of the largest score in each row of X @ W + bias. When two scores tie, the smaller index wins.

This lesson needs PyTorch, which does not run in the browser. The app gives you the exercise as a file with its tests, to run on your own computer.

Input
X = [[1, 2], [3, 0]]
W = [[1, 0, 2], [0, 1, 1]]
bias = [0, 0, -4]

Output
[1, 0]

Input
X = [[1, 0], [0, 1]]
W = [[3, 1], [1, 4]]
bias = [0, 0]

Output
[0, 1]

Input
X = [[0]]
W = [[0, 0]]
bias = [1, 1]

Output
[0]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   Three tensors in, one line of arithmetic, one call to find the winners.

2. Hint 2: The key step

   The product of the two tables is written X @ W. Add the bias to it, then call .argmax(dim=1) to get one index per row.

3. Hint 3: The outline

   Make float64 tensors from X, W and bias. Set logits = x @ w + b. Return logits.argmax(dim=1).tolist().

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. One score per class for every example. [examples, features] @ [features, classes] gives [examples, classes]. The bias has one number per class and is added to every row.
2. The class with the largest score in each row. dim=1 runs across the classes, so each example gets one index.

The same numbered steps are the comments in the solution.

import torch


def solve(X, W, bias):
    x = torch.tensor(X, dtype=torch.float64)
    w = torch.tensor(W, dtype=torch.float64)
    b = torch.tensor(bias, dtype=torch.float64)
    # 1. One score per class for every example
    logits = x @ w + b
    # 2. The class with the largest score in each row
    return logits.argmax(dim=1).tolist()

Watch out

• There is no need for softmax here. Softmax keeps the order of the scores, so the largest score and the largest probability are in the same place.
• argmax(dim=0) would compare examples with each other and return one index per class.

Check yourself