Problem · Array

Streaming Entropy, Part 2: Numerically Stable Entropy

Learn this problem
MediumOpenAIFULLTIMEONSITE INTERVIEW

Problem statement

Streaming Entropy Interview Sequence

This is Part 2 of one four-part ML coding exercise.

As in Part 1, logits define a softmax distribution:

p[i] = exp(x[i]) / sum(exp(x[j]))

and its entropy is

H = -sum(p[i] * ln(p[i]))

This time the logits may be extremely large or small. Directly computing exp(x[i]) can overflow, and computing log(exp(x[i]) / sumExp) can take the logarithm of a probability that rounded to zero.

Numerical Stability Requirements

Return the non-negative entropy without overflow, underflow-driven NaN, or loss of the log-softmax term.

Use a maximum shift for the exponentials and derive log-softmax algebraically from the shifted logits and their normalization sum. Your implementation should run in O(n) time.

An answer is accepted when its absolute or relative error is at most 1e-6.

Continue to Part 3: Block-wise Entropy.

Function

stableEntropy(logits: double[]) → double

Examples

Example 1

logits = [1000.0, 1000.0, -1000.0]return = 0.6931471805599453

Directly evaluating exp(1000) overflows. After shifting by the maximum, the third probability is negligible and the first two are each effectively 0.5.

Example 2

logits = [999999998.0, 999999999.0, 1000000000.0]return = 0.8323955818399389

Subtracting the maximum turns the logits into [-2, -1, 0] without changing the softmax distribution.

Example 3

logits = [42.0]return = 0.0

A one-element softmax distribution is [1], whose entropy is zero.

Constraints

  • 1 <= logits.length <= 2 * 10^5
  • Every logit is finite and lies in [-10^9, 10^9].
  • Use the natural logarithm.
  • Use O(n) time.
  • The result must remain finite for every valid input.

More OpenAI problems

drafts saved locally
public double stableEntropy(double[] logits) {
  // write your code here
}
logits[1000.0, 1000.0, -1000.0]
expected0.6931471805599453
checking account