Optimal transport with entropic regularization
¶

Alfred Galichon (NYU & Sciences Po)

'math+econ+code' masterclass on optimal transport and economic applications
¶

With python code examples
¶

© 2018-2022 by Alfred Galichon. Past and present support from NSF grant DMS-1716489, ERC grant CoG-866274 are acknowledged, as well as inputs from contributors listed here.

If you reuse material from this masterclass, please cite as:
Alfred Galichon, 'math+econ+code' masterclass on optimal transport and economic applications, January 2022. https://github.com/math-econ-code/mec_optim

Learning objectives¶

  • Entropic regularization

  • The log-sum-exp trick

  • Gradient descent, coordinate descent

  • The Iterated Proportional Fitting Procedure (IPFP)

References¶

  • Galichon, Optimal Transport Methods in Economics, Ch. 7.3

  • Peyré, Cuturi, Computational Optimal Transport, Ch. 4.

Entropic regularization of the optimal transport problem¶

Consider the problem

\begin{align*} \max_{\mu\in\mathcal{M}\left( n,m\right) }\sum_{xy}\mu_{xy}\Phi_{xy}-\sigma\sum_{xy}\mu_{xy}\ln\mu_{xy} \end{align*}

where $\sigma>0$. The problem coincides with the optimal assignment problem when $\sigma=0$. When $\sigma\rightarrow+\infty$, the solution to this problem approaches the independent coupling, $\mu_{xy}=n_{x}m_{y}$.

Later on, we will provide microfoundations for this problem, and connect it with a number of important methods in economics (BLP, gravity model, Choo-Siow...). For now, let's just view this as an extension of the optimal transport problem.

We shall compute this problem using Python libraries that we have already met with. Let us start loading them.

In [1]:
# !python -m pip install -i https://pypi.gurobi.com gurobipy ## only if Gurobi not here
In [2]:
import numpy as np
import os
import pandas as pd
import scipy.sparse as spr
import gurobipy as grb
from sklearn import linear_model
from time import time
from scipy.stats import entropy

We define the class OTProblem.

In [3]:
class OTProblem():
    def __init__(self,Φ_x_y,n_x=None,m_y=None):
        self.Φ_a = Φ_x_y.flatten()
        self.nbx , self.nby = Φ_x_y.shape # number of types on each parts
        self.nbz, self.nba = self.nbx + self.nby, self.nbx * self.nby # total number of types and number of arcs
        if n_x is None: 
            self.n_x=np.ones(self.nbx)
        else:
            self.n_x = n_x
        if m_y is None: 
            self.m_y=np.ones(self.nby)
        else:
            self.m_y = m_y    

    def M_z_a(self):
        return spr.vstack([spr.kron(spr.identity(self.nbx), np.ones((1,self.nby))),spr.kron(np.ones((1,self.nbx)),spr.identity(self.nby))])
    
    def solve_full_lp(self,OutputFlag=True):
        m=grb.Model()
        μ_a = m.addMVar(self.nba)
        m.setObjective(self.Φ_a @ μ_a, grb.GRB.MAXIMIZE)
        m.addConstr(self.M_z_a() @ μ_a == np.concatenate([self.n_x,self.m_y]) )
        m.setParam( 'OutputFlag', OutputFlag )
        m.optimize()
        if m.status == grb.GRB.Status.OPTIMAL:
            μopt_a = np.array(m.getAttr('x')).reshape(self.nbx,self.nby)
            popt_z = m.getAttr('pi')
            uopt_x, vopt_y = popt_z[:self.nbx],popt_z[self.nbx:]
        return (μopt_a.reshape((self.nbx,-1)),uopt_x,vopt_y)

Note that in the above, Gurobi is for benchmark purposes with the case $\sigma=0$, but is not suited to compute the nonlinear optimization problem above.

Now, let's load up the affinitymatrix.csv, Xvals.csv and Yvals.csv that you will recall from the previous block. We will work on a smaller population, with nbx types of men and nby types of women.

In [4]:
thepath = 'https://raw.githubusercontent.com/math-econ-code/mec_optim_2021-01/master/data_mec_optim/marriage_personality-traits/'
data_X,data_Y = pd.read_csv(thepath + "Xvals.csv"), pd.read_csv(thepath + "Yvals.csv")
sdX,sdY = data_X.std().values, data_Y.std().values
mX,mY = data_X.mean().values, data_Y.mean().values
feats_x_k, feats_y_l = ((data_X-mX)/sdX).values, ((data_Y-mY)/sdY).values
nbx,nbk = feats_x_k.shape
nby,nbl = feats_y_l.shape
A_k_l = pd.read_csv(thepath + "affinitymatrix.csv").iloc[0:nbk,1:nbl+1].values
Φ_x_y = feats_x_k @ A_k_l @ feats_y_l.T

We extract a smaller matrix of size 50x30 to speed up our numerical explorations.

In [5]:
nbx,nby = 50,30
marriage_ex = OTProblem(Φ_x_y[:nbx,:nby],np.ones(nbx) / nbx, np.ones(nby) / nby)
nrow , ncol = min(8, nbx) , min(8, nby) # number of rows / cols to displayc

As a warm-up, let us compute as in the previous lecture the solution to the problem for $\sigma=0$ that we can compute with Gurobi.

In [6]:
def LPsolve(self):
    ptm = time()
    μ_x_y,u_x,v_y = marriage_ex.solve_full_lp(OutputFlag = False)
    taken = time() - ptm 
    valobs = self.Φ_a.dot(μ_x_y.flatten())
    valtot = valobs
    ite = None
    return μ_x_y,u_x,v_y,valobs,valtot,ite, taken,'LP via Gurobi'

OTProblem.LPsolve = LPsolve

The following function will display the output and allow for a benchmark:

In [7]:
def display( args ):
    μ_x_y,u_x,v_y,valobs,valtot,iterations, taken,name = args
    print('*'*60)
    print('*'* (30 -len(name) // 2) + '  '+name + '  ' + '*'*(26 - (1+len(name)) // 2 ) )
    print('*'*60)
    print('Converged in ', iterations, ' steps and ', taken, 's.')
    print('Sum(mu*Phi)+σ*Sum(mu*log(mu))= ', valtot)
    print('Sum(mu*Phi)                = ', valobs)
    print('*'*60)
    return
In [8]:
display(marriage_ex.LPsolve())

Set parameter WLSAccessID
Set parameter WLSSecret
Set parameter LicenseID
Academic license - for non-commercial use only - registered to alfred.galichon@sciences-po.fr
Academic license - for non-commercial use only - registered to alfred.galichon@sciences-po.fr
************************************************************
************************  LP via Gurobi  *******************
************************************************************
Converged in  None  steps and  1.1296052932739258 s.
Sum(mu*Phi)+σ*Sum(mu*log(mu))=  1.079557261953547
Sum(mu*Phi)                =  1.079557261953547
************************************************************

Dual of the regularized problem¶

Let's compute the dual by the minimax approach. We have

\begin{align*} \max_{\mu\geq0}\min_{u,v}\sum_{xy}\mu_{xy}\left( \Phi_{xy}-u_{x}-v_{y}% -\sigma\ln\mu_{xy}\right) +\sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}% \end{align*}

thus

\begin{align*} \min_{u,v}\sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}+\max_{\mu\geq0}\sum_{xy}% \mu_{xy}\left( \Phi_{xy}-u_{x}-v_{y}-\sigma\ln\mu_{xy}\right) \end{align*}

By FOC in the inner problem, one has $\Phi_{xy}-u_{x}-v_{y}-\sigma\ln \mu_{xy}-\sigma=0,$thus

\begin{align*} \mu_{xy}=\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}-\sigma}{\sigma}\right) \end{align*}

and $\mu_{xy}\left( \Phi_{xy}-u_{x}-v_{y}-\sigma\ln\mu_{xy}\right) =\sigma\mu_{xy}$, thus the dual problem is

\begin{align*} \min_{u,v}\sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}+\sigma\sum_{xy}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}-\sigma}{\sigma}\right) . \end{align*}

After replacing $v_{y}$ by $v_{y}+\sigma$, the dual is

\begin{align*} \min_{u,v}\sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}+\sigma\sum_{xy}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) -\sigma. \tag{V1} \end{align*}

Another expression of the dual¶

Claim: the problem is equivalent to

\begin{align*} \min_{u,v}\sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}+\sigma\log\sum_{i,j} \exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) \tag{V2} \end{align*}

Indeed, let us go back to the minimax expression

\begin{align*} \min_{u,v}\sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}+\max_{\mu\geq0}\sum_{xy}\mu_{xy}\left( \Phi_{xy}-u_{x}-v_{y}-\sigma\ln\mu_{xy}\right) \end{align*}

we see that the solution $\mu$ has automatically $\sum_{xy}\mu_{xy}=1$; thus we can incorporate the constraint into

\begin{align*} \min_{u,v}\sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}+\max_{\mu\geq0:\sum_{xy}\mu_{xy}=1}\sum_{xy}\mu_{xy}\left( \Phi_{xy}-u_{x}-v_{y}-\sigma\ln\mu_{xy}\right) \end{align*}

which yields the our desired result.

This expression is interesting because, taking any $\hat{\mu}\in M\left( n,m\right)$, it reexpresses as

\begin{align*} \max_{u,v}\sum_{xy}\hat{\mu}_{xy}\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) -\log\sum_{xy}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) \end{align*}

therefore if the parameter is $\theta=\left( u,v\right)$, observations are $xy$ pairs, and the likelihood of $xy$ is

\begin{align*} \mu_{xy}^{\theta}=\frac{\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma }\right) }{\sum_{xy}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) } \end{align*}

Hence, our expression will coincide with the maximum likelihood in this model.

A third expression of the dual problem¶

Consider

\begin{align*} \min_{u,v} & \sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y} \\ s.t. \quad & \sum_{i,j}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) =1 \end{align*}

It is easy to see that the solutions of this problem coincide with version 2. Indeed, the Lagrange multiplier is forced to be one. In other words,

\begin{align*} \min_{u,v} & \sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}\\ s.t. \quad & \sigma\log\sum_{i,j}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma }\right) =0 \end{align*}

Small-temperature limit and the log-sum-exp trick¶

Recall that when $\sigma\rightarrow0$, one has

\begin{align*} \sigma\log\left( e^{a/\sigma}+e^{b/\sigma}\right) \rightarrow\max\left( a,b\right) \end{align*}

Indeed, letting $m=\max\left( a,b\right)$,

\begin{align*} \sigma\log\left( e^{a/\sigma}+e^{b/\sigma}\right) =m+\sigma\log\left(\exp\left( \frac{a-m}{\sigma}\right) +\exp\left( \frac{b-m}{\sigma}\right)\right), \end{align*} and the argument of the logarithm lies between $1$ and $2$.

This simple remark is actually a useful numerical recipe called the log-sum-exp trick: when $\sigma$ is small, using the formula above to compute $\sigma\log\left( e^{a/\sigma}+e^{b/\sigma}\right)$ ensures the exponentials won't blow up.

The log-sum-exp trick for regularized OT¶

Back to the third expression, with $\sigma\rightarrow0$, one has

\begin{align*} \min_{u,v} & \sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}\tag{V3}\\ s.t. & \max_{xy}\left( \Phi_{xy}-u_{x}-v_{y}\right) =0\nonumber \end{align*}

This is exactly equivalent with the classical Monge-Kantorovich expression

\begin{align*} \min_{u,v} & \sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}\tag{V3}\\ s.t. & \Phi_{xy}-u_{x}-v_{y}\leq0\nonumber \end{align*}

Back to the third expression of the dual, with $\sigma\rightarrow0$, one has

\begin{align*} \min_{u,v} & \sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}\tag{V3}\\ s.t. & \max_{xy}\left( \Phi_{xy}-u_{x}-v_{y}\right) =0\nonumber \end{align*}

This is exactly equivalent with the classical Monge-Kantorovich expression

\begin{align*} \min_{u,v} & \sum_{x}u_{x}n_{x}+\sum_{y}v_{y}m_{y}\tag{V3}\\ s.t. & \Phi_{xy}-u_{x}-v_{y}\leq0\nonumber \end{align*}

Computation¶

We can compute $\min F\left( x\right)$ by two methods:

Either by gradient descent: $x\left( t+1\right) =x_{t}-\epsilon _{t}\nabla F\left( x_{t}\right) $. (Steepest descent has $\epsilon _{t}=1/\left\vert \nabla F\left( x_{t}\right) \right\vert $.)

Or by coordinate descent: $x_{k}\left( t+1\right) =\arg\min_{x_{k}}F\left( x_{k},x_{-k}\left( t\right) \right)$.

Why do these methods converge? Let's provide some justification. We will decrease $x_{t}$ by $\epsilon d_{t}$, were $d_{t}$ is normalized by $\left\vert d_{t}\right\vert _{p}:=\left( \sum_{i=1}^{n}d_{t}^{i}\right) ^{1/p}=1$. At first order, we have

\begin{align*} F\left( x(t)-\epsilon d_{t}\right) =F\left( x(t)\right) -\epsilon d_{t}^{\intercal}\nabla F\left( x (t)\right) +O\left( \epsilon^{1}\right). \end{align*}

We need to maximize $d_{t}^{\intercal}\nabla F\left( x(t)\right)$ over $\left\vert d_{t}\right\vert _{p}=1$.

  • For $p=2$, we get $d_{t}=\nabla F\left( x(t)\right) /\left\vert \nabla F\left( x(t)\right) \right\vert $

  • For $p=1$, we get $d_{t}=sign\left( \partial F\left( x(t)\right)/\partial x_{k}\right) $ if $\left\vert \partial F\left( x_{t}\right) /\partial x_{k}\right\vert =\max_{l}\left\vert \partial F\left( x(t)\right) /\partial x_{l}\right\vert $, $0$ otherwise.

In our context, gradient descent is

\begin{align*} u_{x}\left( t+1\right) & =u_{x}\left( t\right) -\epsilon\frac{\partial F}{\partial u_{x}}\left( u\left( t\right) ,v\left( t\right) \right) ,\text{ and }\\ v_{y}\left( t+1\right) & =v_{y}\left( t\right) -\epsilon\frac{\partial F}{\partial v_{y}}\left( u\left( t\right) ,v\left( t\right) \right) \end{align*}

while coordinate descent is

\begin{align*} \frac{\partial F}{\partial u_{x}}\left( u_{x}\left( t+1\right) ,u_{-i}\left( t\right) ,v\left( t\right) \right) =0,\text{ and } \frac{\partial F}{\partial v_{y}}\left( u\left( t\right) ,v_{y}\left( t+1\right) ,v_{-j}\left( t\right) \right) =0. \end{align*}

Gradient descent¶

Gradient of objective function in version 1 of our problem:

\begin{align*} \left( n_{x}-\sum_{y}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) ,m_{y}-\sum_{x}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) \right) \end{align*}

Gradient of objective function in version 2

\begin{align*} \left( n_{x}-\frac{\sum_{y}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma }\right) }{\sum_{xy}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) },m_{y}-\frac{\sum_{x}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) }{\sum_{xy}\exp\left( \frac{\Phi_{xy}-u_{x}-v_{y}}{\sigma}\right) }\right) \end{align*}

Coordinate descent¶

Coordinate descent on objective function in version 1:

\begin{align*} n_{x} & =\sum_{y}\exp\left( \frac{\Phi_{xy}-u_{x}\left( t+1\right) -v_{y}\left( t\right) }{\sigma}\right) ,\\ m_{y} & =\sum_{x}\exp\left( \frac{\Phi_{xy}-u_{x}\left( t\right) -v_{y}\left( t+1\right) }{\sigma}\right) \end{align*}

that is

\begin{align*} \left\{ \begin{array} [c]{c} u_{x}\left( t+1\right) =\sigma\log\left( \frac{1}{n_{x}}\sum_{y}\exp\left( \frac{\Phi_{xy}-v_{y}\left( t\right) }{\sigma}\right) \right) \\ v_{y}\left( t+1\right) =\sigma\log\left( \frac{1}{m_{y}}\sum_{x}\exp\left( \frac{\Phi_{xy}-u_{x}\left( t\right) }{\sigma}\right) \right) \end{array} \right. \end{align*}

this is called the Iterated Fitting Proportional Procedure (IPFP), or Sinkhorn's algorithm.

Coordinate descent on objective function in version 2 does not yield a closed-form expression.

IPFP, matrix version¶

Letting $a_{x}=\exp\left( -u_{x}/\sigma\right) $ and $b_{y}=\exp\left( -v_{y}/\sigma\right) $ and $K_{xy}=\exp\left( \Phi_{xy}/\sigma\right) $, one has $\mu_{xy}=a_{x}b_{y}K_{xy}$, and the procedure reexpresses as

\begin{align*} \left\{ \begin{array} [c]{l}% a_{x}\left( t+1\right) =n_{x}/\left( Kb\left( t\right) \right) _{x}\text{ and }\\ b_{y}\left( t+1\right) =m_{y}/\left( K^{\intercal}a\left( t\right) \right) _{y}. \end{array} \right. \end{align*}

Because this algorithm involves matrix operations only, and is naturally suited for parallel computation, GPUs are a tool of choice for addressing is. See chap. 4 of Peyré and Cuturi.

Implementation¶

First, as a convenience, we would like to have a function that computes $\sum_k \mu_k \log \mu_k$ without failing when some of the entries of $\mu_k$ are zero. We code this into:

In [9]:
def sum_xlogx(a):
    s=a.sum()
    return s*np.log(s) - s * entropy(a.flatten(),axis=None)

Let's implement this algorithm. Return to the matrix-IPFP algorithm:

In [10]:
def matrixIPFP(self,σ , tol = 1e-9, maxite = 1e+06 ):
    ptm = time()
    ite = 0
    K_x_y = np.exp(self.Φ_a / σ).reshape(self.nbx,-1)
    B_y = np.ones(self.nby)
    error = tol + 1
    while error > tol and ite < maxite:
        A_x = self.n_x / (K_x_y @ B_y)
        KA_y = (A_x.T @ K_x_y)
        error = (abs(KA_y * B_y / self.m_y)-1).max()
        B_y = self.m_y / KA_y
        ite = ite + 1
        
    u_x,v_y = - σ * np.log(A_x),- σ * np.log(B_y)
    μ_x_y = K_x_y * A_x.reshape((-1,1)) * B_y.reshape((1,-1))
    valobs = self.Φ_a.dot(μ_x_y.flatten())
    valtot =  valobs - σ * sum_xlogx(μ_x_y)
    taken = time() - ptm
    if ite >= maxite:
        print('Maximum number of iteations reached in matrix IPFP.')    
    return μ_x_y,u_x,v_y,valobs,valtot,ite, taken, 'matrix IPFP'

OTProblem.matrixIPFP = matrixIPFP
In [11]:
display(marriage_ex.matrixIPFP(0.1))

************************************************************
*************************  matrix IPFP  ********************
************************************************************
Converged in  229  steps and  0.002029895782470703 s.
Sum(mu*Phi)+σ*Sum(mu*log(mu))=  1.5590368681937061
Sum(mu*Phi)                =  1.0283148199160483
************************************************************

To see the benefit of the matrix version, let us recode the same algorithm as above, but in the log-domain, namely iterate over the values of $u$ and $v$.

In [12]:
def logdomainIPFP(self,σ = 0.1, tol = 1e-9, maxite = 1e+06 ):
    ptm = time()
    ite = 0
    Φ_x_y = self.Φ_a.reshape(self.nbx,-1)
    v_y = np.zeros(self.nby)
    λ_x,ζ_y = - σ * np.log(self.n_x), - σ * np.log(self.m_y)
    error = tol + 1
    while error > tol and ite < maxite:
        u_x = λ_x + σ * np.log( (np.exp((Φ_x_y - v_y.reshape((1,-1)))/σ)).sum( axis=1) )
        KA_y = (np.exp((Φ_x_y -u_x.reshape((-1,1))) / σ)).sum(axis=0)
        error = np.max(np.abs(KA_y * np.exp(-v_y / σ) / self.m_y - 1))
        v_y = ζ_y + σ * np.log(KA_y)
        ite = ite + 1
    
    μ_x_y =np.exp((Φ_x_y -u_x.reshape((-1,1)) - v_y.reshape((1,-1)))/σ )
    valobs = self.Φ_a.dot(μ_x_y.flatten())
    valtot =  valobs - σ * sum_xlogx(μ_x_y)
    taken = time() - ptm
    if ite >= maxite:
        print('Maximum number of iteations reached in log-domain IPFP.')
    return μ_x_y,u_x,v_y,valobs,valtot,ite, taken, 'log-domain IPFP'

OTProblem.logdomainIPFP = logdomainIPFP
In [13]:
display(marriage_ex.logdomainIPFP(0.1))
display(marriage_ex.matrixIPFP(0.1))

************************************************************
***********************  log-domain IPFP  ******************
************************************************************
Converged in  229  steps and  0.008599281311035156 s.
Sum(mu*Phi)+σ*Sum(mu*log(mu))=  1.5590368681937066
Sum(mu*Phi)                =  1.0283148199160488
************************************************************
************************************************************
*************************  matrix IPFP  ********************
************************************************************
Converged in  229  steps and  0.0017642974853515625 s.
Sum(mu*Phi)+σ*Sum(mu*log(mu))=  1.5590368681937061
Sum(mu*Phi)                =  1.0283148199160483
************************************************************

We see that the log-domain IPFP, while mathematically equivalent to matrix IPFP, it is noticeably slower.

IPFP with the log-sum-exp trick¶

The matrix IPFPis very fast, partly due to the fact that it involves linear algebra operations. However, it breaks down when $\sigma$ is small; this is best seen taking a log transform and returning to $u^{k}=-\sigma\log a^{k}$ and $v^{k}=-\sigma\log b^{k}$, that is

\begin{align*} \left\{ \begin{array} [c]{l}% u_{x}^{k}=\mu_{x}+\sigma\log\sum_{y}\exp\left( \frac{\Phi_{xy}-v_{y}^{k-1}% }{\sigma}\right) \\ v_{y}^{k}=\zeta_{y}+\sigma\log\sum_{x}\exp\left( \frac{\Phi_{xy}-u_{x}^{k}% }{\sigma}\right) \end{array} \right. \end{align*}

where $\mu_{x}=-\sigma\log n_{x}$ and $\zeta_{y}=-\sigma\log m_{y}$.

One sees what may go wrong: if $\Phi_{xy}-v_{y}^{k-1}$ is positive in the exponential in the first sum, then the exponential blows up due to the small $\sigma$ at the denominator. However, the log-sum-exp trick can be used in order to avoid this issue.

Consider

\begin{align*} \left\{ \begin{array} [c]{l}% \tilde{v}_{x}^{k}=\max_{y}\left\{ \Phi_{xy}-v_{y}^{k}\right\} \\ \tilde{u}_{y}^{k}=\max_{x}\left\{ \Phi_{xy}-u_{x}^{k}\right\} \end{array} \right. \end{align*}

(the indexing is not a typo: $\tilde{v}$ is indexed by $i$ and $\tilde{u}$ by $j$).

One has

\begin{align*} \left\{ \begin{array} [c]{l}% u_{x}^{k}=\mu_{x}+\tilde{v}_{x}^{k-1}+\sigma\log\sum_{y}\exp\left( \frac {\Phi_{xy}-v_{y}^{k-1}-\tilde{v}_{x}^{k}}{\sigma}\right) \\ v_{y}^{k}=\zeta_{y}+\tilde{u}_{y}^{k}+\sigma\log\sum_{x}\exp\left( \frac {\Phi_{xy}-u_{x}^{k}-\tilde{u}_{y}^{k}}{\sigma}\right) \end{array} \right. \end{align*}

and now the arguments of the exponentials are always nonpositive, ensuring the exponentials don't blow up.

Both the matrix version and the log-domain version of the IPFP will break down when $\sigma$ is small, e.g. $\sigma=0.001$ (Try!). However if we modify the second procedure using the log-sum-exp trick, things work again:

In [14]:
def logdomainIPFP_with_LSE_trick(self,σ , tol = 1e-9, maxite = 1e+06 ):
    ptm = time()
    ite = 0
    Φ_x_y = self.Φ_a.reshape(self.nbx,-1)
    v_y = np.zeros(self.nby)
    λ_x,ζ_y = - σ * np.log(self.n_x), - σ * np.log(self.m_y)
    error = tol + 1
    while error > tol and ite < maxite:
        vstar_x = (Φ_x_y - v_y.reshape((1,-1))).max( axis = 1)
        u_x = λ_x + vstar_x + σ * np.log( (np.exp((Φ_x_y - vstar_x.reshape((-1,1)) - v_y.reshape((1,-1)))/σ)).sum( axis=1) )
        ustar_y = (Φ_x_y - u_x.reshape((-1,1)) ).max( axis = 0)
        KA_y = (np.exp((Φ_x_y -u_x.reshape((-1,1)) - ustar_y.reshape((1,-1)) ) / σ)).sum(axis=0)
        error = np.max(np.abs(KA_y * np.exp( (ustar_y-v_y) / σ) / self.m_y - 1))
        v_y = ζ_y + ustar_y+ σ * np.log(KA_y)
        ite = ite + 1
    μ_x_y =np.exp((Φ_x_y -u_x.reshape((-1,1)) - v_y.reshape((1,-1)))/σ )
    valobs = self.Φ_a.dot(μ_x_y.flatten())
    valtot =  valobs - σ * sum_xlogx(μ_x_y)
    taken = time() - ptm
    if ite >= maxite:
        print('Maximum number of iteations reached in log-domain IPFP with LSE trick.')
    return μ_x_y,u_x,v_y,valobs,valtot,ite, taken, 'log-domain IPFP with LSE trick'

OTProblem.logdomainIPFP_with_LSE_trick = logdomainIPFP_with_LSE_trick
In [15]:
display(marriage_ex.logdomainIPFP_with_LSE_trick(0.1))
display(marriage_ex.logdomainIPFP(0.1))

************************************************************
***************  log-domain IPFP with LSE trick  ***********
************************************************************
Converged in  229  steps and  0.012583017349243164 s.
Sum(mu*Phi)+σ*Sum(mu*log(mu))=  1.559036868193707
Sum(mu*Phi)                =  1.0283148199160492
************************************************************
************************************************************
***********************  log-domain IPFP  ******************
************************************************************
Converged in  229  steps and  0.00818490982055664 s.
Sum(mu*Phi)+σ*Sum(mu*log(mu))=  1.5590368681937066
Sum(mu*Phi)                =  1.0283148199160488
************************************************************

In contrast, when $\sigma = 0.01$ we see that the algorithm works with the log-sum-exp trick but fails without:

In [16]:
# display(marriage_ex.logdomainIPFP_with_LSE_trick(0.1))
# display(marriage_ex.logdomainIPFP(0.001))

Computations using GLM¶

Recall that the margining-out matrix of shape $\left( nbx+nby\right) \times \left( nbx.nby\right) $ is given by

$M=\binom{I_{X}\otimes 1_{Y}^{\top }}{1_{Y}^{\top }\otimes I_{Y}}$

Introduce $\hat{\mu}=nm^{\top } / (\sum_x n_x)$

We have $M\hat{\mu}=\binom{n}{m}$

The problem writes

$\min_{p}1_{\mathcal{A}}^{\top }\exp \left( \frac{\Phi -M^{\top }p}{\sigma }% \right) -\hat{\mu}^{\top }\left( \frac{\Phi -M^{\top }p}{\sigma }\right) $

And setting $\tilde{p}=p/\sigma $ and $\tilde{\Phi}=\Phi /\sigma $ yields that $\tilde{p}$ is obtained by

$\min_{\tilde{p}}1^{\top }\exp \left( \tilde{\Phi}-M^{\top }\tilde{p}\right) -\hat{\mu}^{\top }\left( \tilde{\Phi}-M^{\top }\tilde{p}\right) $

which we can rewrite as a weighted Poisson regression

$\min_{\tilde{p}}\sum_{a\in \mathcal{A}}w_{a}\exp \left( -\left( M^{\top }% \tilde{p}\right) _{a}\right) -\sum_{a\in \mathcal{A}}w_{a}\hat{\mu}_{a}e^{-% \tilde{\Phi}_{a}}\left( \tilde{\Phi}-M^{\top }\tilde{p}\right) _{a}$

where $w_{a}=\exp \tilde{\Phi}_{a}$

Dropping the constant term, this implements

$\min_{\tilde{p}}\sum_{a\in \mathcal{A}}w_{a}\exp \left( -M^{\top }\tilde{p}% \right) _{a}+\sum_{a\in \mathcal{A}}w_{a}\left( \hat{\mu}e^{-\tilde{\Phi}% }\right) _{a}\left( M^{\top }\tilde{p}\right) _{a}$

In [17]:
def solveGLM(self,σ , tol = 1e-9):
    ptm = time()
    muhat_a = (self.n_x.reshape((self.nbx,-1)) @ self.m_y.reshape((-1,self.nby))).flatten() / self.n_x.sum()
    ot_as_glm = linear_model.PoissonRegressor(fit_intercept=False,tol=tol ,verbose=3,alpha=0)
    ot_as_glm.fit( - self.M_z_a().T, muhat_a * np.exp(-self.Φ_a / σ) , sample_weight = np.exp(self.Φ_a / σ))

    p = σ * ot_as_glm.coef_
    u_x,v_y  = p[:self.nbx] - p[0], p[self.nbx:]+p[0]
    μ_x_y =np.exp((self.Φ_a.reshape((self.nbx,-1)) -u_x.reshape((-1,1)) - v_y.reshape((1,-1)))/σ )
    valobs = self.Φ_a.dot(μ_x_y.flatten())
    valtot =  valobs - σ * sum_xlogx(μ_x_y)
    taken = time() - ptm
    return μ_x_y,u_x,v_y,valobs,valtot,None, taken, 'GLM'

OTProblem.solveGLM = solveGLM
In [18]:
σ = 0.5
display(marriage_ex.solveGLM(σ ))
display(marriage_ex.matrixIPFP(σ ))

RUNNING THE L-BFGS-B CODE

           * * *

Machine precision = 2.220D-16
 N =           80     M =           10

           * * *

Tit   = total number of iterations
Tnf   = total number of function evaluations
Tnint = total number of segments explored during Cauchy searches
Skip  = number of BFGS updates skipped
Nact  = number of active bounds at final generalized Cauchy point
Projg = norm of the final projected gradient
F     = final function value

           * * *

   N    Tit     Tnf  Tnint  Skip  Nact     Projg        F
************************************************************
*****************************  GLM  ************************
************************************************************
Converged in  None  steps and  0.02057504653930664 s.   80     25     27      1     0     0   5.535D-10   1.046D-04

CONVERGENCE: NORM_OF_PROJECTED_GRADIENT_<=_PGTOL            

Sum(mu*Phi)+σ*Sum(mu*log(mu))=  4.0611956293694105
Sum(mu*Phi)                =  0.6764482965501332
************************************************************
************************************************************
*************************  matrix IPFP  ********************
************************************************************
Converged in  28  steps and  0.0012094974517822266 s.
Sum(mu*Phi)+σ*Sum(mu*log(mu))=  4.061144010352285
Sum(mu*Phi)                =  0.6764437848144771
************************************************************

 This problem is unconstrained.

However, when $\sigma = 0.1$ the GLM approach fails:

In [19]:
σ = 0.1
display(marriage_ex.solveGLM(σ ))
display(marriage_ex.matrixIPFP(σ ))

RUNNING THE L-BFGS-B CODE

           * * *

Machine precision = 2.220D-16
 N =           80     M =           10

           * * *

Tit   = total number of iterations
Tnf   = total number of function evaluations
Tnint = total number of segments explored during Cauchy searches
Skip  = number of BFGS updates skipped
Nact  = number of active bounds at final generalized Cauchy point
Projg = norm of the final projected gradient
F     = final function value

           * * *

   N    Tit     Tnf  Tnint  Skip  Nact     Projg        F
   80     29     30      1     0     0   6.696D-10   4.326D-09

CONVERGENCE: NORM_OF_PROJECTED_GRADIENT_<=_PGTOL            
************************************************************
*****************************  GLM  ************************
************************************************************
Converged in  None  steps and  0.013177633285522461 s.
Sum(mu*Phi)+σ*Sum(mu*log(mu))=  2748911.438926413
Sum(mu*Phi)                =  5722389.281066181
************************************************************
************************************************************
*************************  matrix IPFP  ********************
************************************************************
Converged in  229  steps and  0.0037682056427001953 s.
Sum(mu*Phi)+σ*Sum(mu*log(mu))=  1.5590368681937061
Sum(mu*Phi)                =  1.0283148199160483
************************************************************

 This problem is unconstrained.