1 / 1

Value and Grad of a Loss

JAX

Easy

Compute the loss of a small logistic regression model and its gradient with respect to every parameter.

The model has a weight vector w and a scalar bias b, stored in the dict params under the keys w and b. For the n-th example, with feature row $x_n$ and label $y_n \in \{0, 1\}$ , the logit is $z_n = x_n \cdot w + b$ , and the loss is

L(w, b) = \frac{1}{N} \sum_{n=1}^{N} \Big( \log\big(1 + e^{z_n}\big) - y_n z_n \Big) + \frac{\lambda}{2} \sum_{i=1}^{D} w_i^2

Here x has shape (N, D) with rows $x_n$ , y has shape (N,), N is the number of examples, D is the number of features, and $\lambda$ is the argument l2. The penalty covers w only and never b.

Return a dict with two keys. The value under loss is $L$ as a float32 JAX array of shape (). The value under grads is a dict with keys w and b holding $\partial L / \partial w$ as a float32 JAX array of shape (D,) and $\partial L / \partial b$ as a float32 JAX array of shape ().

Examples

Input: params = {"w":[0.5,-1.0],"b":0.25}, x = [[1.0,2.0],[-1.0,0.5],[0.0,-1.0],[2.0,1.0]], y = [0.0,1.0,1.0,0.0], l2 = 0.1

Output: {"loss":0.679167,"grads":{"w":[0.556558,0.122672],"b":-0.029251}}

Explanation: The logits are x @ w + 0.25 = [-1.25, -0.75, 1.25, 0.25]. The mean logistic loss over the four examples is 0.616667, and the penalty adds 0.1 / 2 * (0.5^2 + 1^2) = 0.0625, giving 0.679167. The bias gradient has no penalty term.

Input: params = {"w":[0.2,0.0,-0.4],"b":-0.5}, x = [[1.0,0.0,2.0],[0.5,1.5,-1.0],[-2.0,1.0,0.0]], y = [1.0,0.0,1.0], l2 = 0.0

Output: {"loss":1.107212,"grads":{"w":[0.307213,0.013017,-0.66684],"b":-0.320403}}

Input: params = {"w":[3.0],"b":1.0}, x = [[40.0],[-40.0]], y = [0.0,1.0], l2 = 0.01

Output: {"loss":120.045,"grads":{"w":[40.03],"b":0.0}}

Hint 1

loss, grads = jax.value_and_grad(loss_fn)(params) returns both in one call.

Hint 2

jax.nn.softplus(z) computes log(1 + exp(z)) without overflow.

Hint 3

grads is a dict with the same keys as params.

Requirements

Write the loss as a function of the params dict and differentiate it with jax.value_and_grad, which returns the loss and the gradient dict from one forward and backward pass. Solutions that never call jax.value_and_grad are rejected.
Evaluate log(1 + e^z) in a form that stays finite when z is larger than about 88.7, where e^z overflows float32.
Leave the params dict unchanged.

Constraints

params holds w, a float32 JAX array of shape (D,), and b, a float32 JAX array of shape ().
x is a float32 JAX array of shape (N, D) with N between 1 and 8 and D between 1 and 4.
y is a float32 JAX array of shape (N,) whose entries are 0 or 1.
l2 is a non-negative Python float.
Logits can reach about 125 in absolute value.

Try Similar Problems

Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Functional Array UpdateEasy PRNG Key SplittingEasy Masked Mean Without Boolean IndexingEasy Static Arguments in jitEasy Model Ensemble with in_axesEasy SGD with Momentum over a PytreeEasy Rotary Position EmbeddingEasy KV Cache Update at a PositionEasy Device Mesh and Named ShardingEasy