Output Layer — Module 6.1

How This Lesson Fits Module 6.1

The output layer converts learned representations from hidden layers into task-specific predictions. It is the point where architecture, target encoding, activation choice, and loss function must agree.

Learning Objectives

By the end of this lesson, students should be able to:

Define the output layer as the model's final prediction interface.
Match output dimensions to regression, binary classification, and multiclass classification.
Explain logits and why many PyTorch losses expect raw scores.
Choose suitable final activations for inference.
Implement output heads for common tasks in PyTorch.
Diagnose target-output mismatches.

Definition

The output layer is the final network layer that maps hidden representations into values interpreted as predictions, logits, probabilities, or scores for a specific task.

Task Determines the Head

A model body can learn useful features, but the output head makes those features usable. Regression may need one linear value. Binary classification commonly emits one logit. Multiclass classification emits one logit per class, later interpreted by softmax. In PyTorch, keep the distinction clear: training losses often expect raw logits, while user-facing inference may convert logits to probabilities.

Task	Output shape	Typical PyTorch loss
Regression	`[N, 1]` or `[N]`	`nn.MSELoss` / `nn.L1Loss`
Binary classification	`[N]` or `[N,1]` logits	`nn.BCEWithLogitsLoss`
Multiclass classification	`[N, C]` logits	`nn.CrossEntropyLoss`
Multilabel classification	`[N, C]` independent logits	`nn.BCEWithLogitsLoss`
Embedding/retrieval	Vector representation	Contrastive or ranking loss

PyTorch Practice

These small heads show how the final layer changes while the hidden representation can remain the same.

import torch
from torch import nn

hidden = torch.randn(8, 32)
regression_head = nn.Linear(32, 1)
binary_head = nn.Linear(32, 1)
multiclass_head = nn.Linear(32, 5)

y_reg = regression_head(hidden)
y_bin_logits = binary_head(hidden).squeeze(1)
y_class_logits = multiclass_head(hidden)

print(y_reg.shape, y_bin_logits.shape, y_class_logits.shape)

Logits, Probabilities, Predictions

Logits

Raw model scores
Can be negative or positive
Preferred by stable PyTorch losses

Probabilities

After sigmoid or softmax
Useful for interpretation
Can saturate numerically if computed too early

Predictions

Threshold or argmax decision
Used for metrics and product behavior
Should match task costs

Strengths and Tradeoffs

Useful because

A correct output head makes training stable and metrics meaningful.
Raw logits preserve numerical stability for common losses.
Task-specific heads allow shared model bodies.

Watch for

Wrong output dimension can silently train the wrong objective.
Applying softmax before CrossEntropyLoss duplicates work and hurts stability.
Thresholds for deployment may need calibration, not just defaults.

How It Flows

1. Represent

Hidden layers produce a feature vector.

2. Project

Output layer maps features to task scores.

3. Train

Loss compares scores with targets.

4. Interpret

Inference converts scores to values, probabilities, or labels.

Common Misconception

A frequent PyTorch mistake is adding Softmax before nn.CrossEntropyLoss. That loss expects raw logits and applies a stable log-softmax internally.

Knowledge Check

Short Answer: What does the output layer produce? Answer: Task-specific predictions, logits, or scores.
True/False: Multiclass classification usually emits one score per class. Answer: True.
Multiple Choice: CrossEntropyLoss expects: (a) raw logits, (b) one-hot strings, (c) already rounded labels as input. Answer: (a).
Short Answer: Binary classification often uses which stable loss? Answer: BCEWithLogitsLoss.
True/False: Output shape should match target encoding. Answer: True.
Short Answer: What is a logit? Answer: A raw unnormalized model score.
Multiple Choice: Regression output is commonly: (a) one continuous value, (b) one score per class, (c) token IDs only. Answer: (a).
Short Answer: How do you turn multiclass logits into a class prediction? Answer: Use argmax over class dimension.
True/False: Deployment thresholds may differ from 0.5. Answer: True.
Short Answer: Why avoid premature softmax in training? Answer: It can reduce numerical stability and duplicate loss behavior.

Key Takeaways

The output layer must match the prediction task and target encoding.
PyTorch training usually prefers raw logits for classification losses.
Inference can convert logits into probabilities or labels after training.
Next, Weights explains the trainable parameters behind every layer.

Trainer’s Guide

Hands-on idea: Give students three target tensors and ask them to choose output dimensions and losses before showing any model code.

Discussion prompt: Why might a product team choose a threshold other than 0.5 for a binary classifier?

Recap: The output layer is the task-specific prediction interface, and it must align with the loss. Continue with Weights.