Neural Mastery
← Practice
Practice: Softmax Gradient (Vector-Jacobian Product)hard

Practice: Softmax Gradient (Vector-Jacobian Product)

hard Neural Networks by Hand 20–25 min

NumPy allowed · plain Python earns 10 bonus points

New to this topic? Deep Learning Overview
Study this concept first →
Interactive Vector Lab

Softmax Activation

Exponentiation and normalisation mapping logits to probabilities: p_i = exp(z_i) / Σ exp(z_j)
Editable Inputs (change values live)Step 0 of 6
z =
Softmax turns raw logits into a normalized probability distribution where Σ p_i = 1.
p[0]: 65.9%
p[1]: 24.2%
p[2]: 9.9%
Live Result[0.659, 0.242, 0.099]

Compute the softmax backpropagation gradient efficiently in O(n) time without constructing the full n x n Jacobian matrix.

Task

Implement softmax_gradient(s, grad_output). Given softmax probabilities s and upstream gradient g, compute the vector-Jacobian product: (J^T g)_i = s_i * (g_i - sum_j s_j * g_j) in O(n) time.

Function Signature

softmax_gradient(s: list[float], grad_output: list[float]) -> list[float]

Examples

Example 1: Uniform distribution with symmetric gradients
Input: s = [0.5, 0.5] grad_output = [1, -1]
Output: [0.5, -0.5]
Example 2: Peaked distribution
Input: s = [0.8, 0.1, 0.1] grad_output = [0.5, -0.2, -0.3]
Output: [0.12, -0.055, -0.065]

Constraints

  • len(s) == len(grad_output) >= 1
  • s is a valid probability distribution
  • O(n) time complexity required
Python3Saved ✓
def softmax_gradient(s: list[float], grad_output: list[float]) -> list[float]:
"""Compute the vector-Jacobian product for softmax in O(n) time.
s is the softmax probability vector, grad_output is dL/ds.
Return dL/dz.
"""
# Your code here
pass

Case 1: Uniform distribution with symmetric gradients
Input: s = [0.5, 0.5] grad_output = [1, -1]
Expected: [0.5, -0.5]
Case 2: Peaked distribution
Input: s = [0.8, 0.1, 0.1] grad_output = [0.5, -0.2, -0.3]
Expected: [0.12, -0.055, -0.065]