Batches, Logits, and Output Heads

Batches are the unit of computation

Networks operate on batches of examples, not single examples. A batch of 32 images is a tensor of shape [32, 3, 224, 224]; a batch of sentences is [32, seq_len, hidden_dim]. The first dimension is almost always the batch dimension and you should reach for it without thinking — x[0] is the first example in the batch.

Batching is what makes GPUs useful. A single forward pass on one example wastes most of the chip's parallelism. Doubling the batch size often halves the per-example time, until memory bandwidth or VRAM caps you.

Tip: If a tutorial uses single-example forward passes, mentally rewrite them as batched. Real training loops always operate in batches; single-example code is a teaching shortcut.

Logits vs probabilities

Logits are the raw, un-normalized outputs of the final linear layer. They can be any real number, positive or negative. Probabilities come from passing logits through softmax (multi-class) or sigmoid (binary). PyTorch's classification losses take logits, not probabilities — they apply softmax internally for numerical stability. Returning logits from your model is the right default.

Output heads

The "head" is the final layer (or small stack) that maps the network's internal representation to the task output. For classification, it is nn.Linear(hidden_dim, num_classes). For regression, nn.Linear(hidden_dim, 1). For multi-label, nn.Linear(hidden_dim, num_labels) with sigmoid + BCE. For embedding learning, an L2-normalized projection.

This is the layer you usually replace when you fine-tune a pretrained model on a new task. Backbone unchanged, head re-initialized for the new label set.

Principle: Return logits, not probabilities. Pair them with the right loss function. Save softmax for visualization or for downstream code that genuinely needs probabilities.

Code

Batched classifier with task-shaped heads·python

import torch, torch.nn as nn

class Backbone(nn.Module):
    def __init__(self, in_dim, hidden_dim):
        super().__init__()
        self.body = nn.Sequential(
            nn.Linear(in_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
        )
    def forward(self, x):
        return self.body(x)

backbone = Backbone(in_dim=512, hidden_dim=256)
head_clf = nn.Linear(256, 10)
head_reg = nn.Linear(256, 1)
head_emb = nn.Linear(256, 64)

x = torch.randn(32, 512)
z = backbone(x)
print(head_clf(z).shape)                         # [32, 10] logits
print(head_reg(z).shape)                         # [32, 1]  predictions
print(nn.functional.normalize(head_emb(z), dim=-1).shape)  # [32, 64]

Batches are the unit of computation

Logits vs probabilities

Output heads

Code

External links

Exercise

Progress

Comments 0