KL Projection Step

Problem

Perform one Newton step of KL projection (moment matching) onto an exponential family. Given probs over values x_values and a target mean, compute the current mean and variance, set theta = (target_mean - mean) / variance, reweight q_i proportional to p_i * exp(theta * x_i), and return the normalized result. Return probs unchanged if the variance is zero or the mean already matches.

Hint

Iterating this step converges to the I-projection of p onto the moment constraint.

Test cases

case 1

in: [[0.5,0.5],[0,1],0.5]

expected: [0.5,0.5]

case 2

in: [[0.5,0.5],[0,1],0.75]

expected: [0.2689414213699951,0.7310585786300049]

case 3

in: [[0.5,0.5],[0,1],0.25]

expected: [0.7310585786300049,0.2689414213699951]

case 4

in: [[0.25,0.5,0.25],[0,1,2],1.5]

expected: [0.07232948812851327,0.39322386648296365,0.534446645388523]

Starter codePython

import math
def kl_projection_step(probs, x_values, target_mean):
    # Your code here
    pass

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