Sinkhorn Divergence Gradient Lite

Problem

The Sinkhorn divergence is S = OT(mu, nu) - 0.5*OT(mu, mu) - 0.5*OT(nu, nu), and by the envelope theorem the gradient with respect to a shared cost entry is the corresponding transport plan entry.

Given the three plans and indices (i, j), return P_munu[i][j] - 0.5*P_mumu[i][j] - 0.5*P_nunu[i][j].

Hint

The self-transport terms correct the entropic bias of Sinkhorn.

Test cases

case 1

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

expected: 0

case 2

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

expected: 0

case 3

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

expected: 0.5

case 4

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

expected: -0.5

Starter codePython

def sinkhorn_divergence_grad(plan_munu, plan_mumu, plan_nunu, i, j):
    # Your code here
    pass

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