Online Softmax Rescale Step

Problem

Perform one step of the online softmax update with a running maximum m and running sum l and a new score x. If x > m, rescale the sum by exp(m - x), add 1 and update m; otherwise add exp(x - m). Return [m_new, l_new].

The signature is online_softmax_step(m, l, x).

Hint

The rescale factor exp(m - x) repairs the accumulated sum when a new maximum appears.

Test cases

case 1

in: [0,1,1]

expected: [1,1.3678794412]

case 2

in: [0,1,-1]

expected: [0,1.3678794412]

case 3

in: [5,0,1]

expected: [5,0.0183156389]

case 4

in: [0,0,0]

expected: [0,1]

case 5

in: [-1000000000,1,0]

expected: [0,1]

Starter codePython

import math
def online_softmax_step(m, l, x):
    # Your code here
    pass

5 test cases · run your code against them in the editor