📘

Apply Softmax for Predictions

Problem Statement

Convert raw model outputs (logits) into probabilities using Softmax.

Background

Neural networks output raw scores called logits. To interpret these as probabilities, you need to apply Softmax, which converts them to values in [0, 1] that sum to 1 across classes.

Your Task

Write a function apply_softmax(logits) that takes a batch of logits and returns statistics about the resulting probability distribution.

Output Format

Return a dictionary with keys: "probabilities_sum" (list of floats, 1 decimal — should each be ~1.0), "predictions" (list of ints — predicted class per sample), "max_probs" (list of floats, 2 decimals — confidence per sample).

Example:

Input:

logits = [[2.0, 1.0, 0.1], [0.5, 2.5, 0.3]]

Output:

{'probabilities_sum': [1.0, 1.0], 'predictions': [0, 1], 'max_probs': [0.66, 0.8]}

Reasoning:

Softmax normalizes to probabilities, argmax gives predicted class

Constraints:

Use nn.Softmax(dim=1)
Round probabilities appropriately
Handle batch dimension correctly

More from Pytorch

Count Model ParametersMedium Build with nn.SequentialEasy Use nn.Flatten LayerEasy Create a Simple Neural Network ClassEasy Compose Multiple TransformsMedium

1. Background Knowledge

Logits are raw, unnormalized scores output by a neural network's final linear layer, ranging from $[-∞, ∞]$ . They represent each class's relative strength but aren't interpretable as probabilities.

The Softmax function converts logits to a probability distribution:

\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^K e^{z_j}}

where $z_i$ is the logit for class $i$ , and $K$ is the number of classes. Key properties:

Outputs in $[0, 1]$
Sums to 1 across classes: $\sum_{i=1}^K$ \text{softmax} $(z_i) = 1$
Higher logits → higher probabilities (exponential scaling amplifies differences)

In PyTorch, nn.Softmax(dim=1) applies this along dimension 1 (class dimension for batch × classes shape), enabling maximum likelihood estimation with cross-entropy loss during training.

Predictions use argmax() on probabilities (or logits, as argmax is invariant to monotonic transforms like softmax).

2. Algorithm Approach

Instantiate Softmax: softmax = nn.Softmax(dim=1)
Apply transformation: probs = softmax(logits) (shape preserved: batch_size × num_classes)
Extract metrics:

Sum: probs.sum(dim=1).tolist() (should be ~1.0 per sample)
Predictions: probs.argmax(dim=1).tolist()
Max probs: probs.max(dim=1).values.tolist()

This is a standard post-processing step in classification pipelines.

3. Step-by-Step Strategy

import torch
import torch.nn as nn

def apply_softmax(logits):
 softmax = nn.Softmax(dim=1)
 probs = softmax(logits)

 probs_sum = probs.sum(dim=1).round(decimals=1).tolist()
 predictions = probs.argmax(dim=1).tolist()
 max_probs = probs.max(dim=1).values.round(decimals=2).tolist()

👑

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

🔒

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.