Sinkhorn balancing and a Sinkhorn balance loss function.

sinkhorn_distributed.py

"""Sinkhorn balancing and a Sinkhorn balance loss function."""

import math
from typing import Optional

import torch
from torch import distributed as dist
from torch.distributed import nn as dnn
from torch.nn import functional as F


def log_sinkhorn(
    L: torch.Tensor,
    *,
    tol: float = 1e-4,
    distributed: bool = False,
    group: Optional[dist.ProcessGroup] = None
) -> torch.Tensor:
    """Balance logits using the Sinkhorn algorithm in log-space.

    Args:
        L: The logits, shape (..., N, M). If `distributed` is `True`, the logits are sharded on
            dimension -2 across ranks (each rank's input is a shard).
        tol: The tolerance for the Sinkhorn iterations.
        distributed: Whether to use distributed Sinkhorn iterations.
        group: The process group to use, if `distributed` is `True`.

    Returns:
        The balanced logits, shape (..., N, M). If `distributed` is `True`, returns the local shard
        of the collectively balanced logits.
    """
    world_size = dist.get_world_size(group=group) if distributed else 1
    a_dual = torch.zeros(*L.shape[:-2], 1, L.shape[-1], device=L.device, dtype=L.dtype)
    b_dual = torch.zeros(*L.shape[:-2], L.shape[-2], 1, device=L.device, dtype=L.dtype)
    offset_a = math.log(L.shape[-1])
    offset_b = math.log(L.shape[-2] * world_size)
    err_prev = torch.full_like(L[..., 0, 0], float("inf"))
    while True:
        a_dual = torch.logsumexp(L - b_dual, dim=-2, keepdim=True) + offset_a
        if distributed:
            a_dual = torch.logsumexp(torch.stack(dnn.all_gather(a_dual, group=group)), dim=0)
        b_dual = torch.logsumexp(L - a_dual, dim=-1, keepdim=True) + offset_b
        with torch.no_grad():
            pi_sum = torch.sum(torch.exp(L - a_dual - b_dual), dim=-2)
            if distributed:
                dist.all_reduce(pi_sum, op=dist.ReduceOp.SUM, group=group)
            err = torch.norm(pi_sum - 1 / L.shape[-1], p=1, dim=-1)
        if torch.all((err <= tol) | (err >= err_prev) | ~err.isfinite()):
            break
        err_prev = err
    return L - a_dual - b_dual


def sinkhorn_balance_loss(
    L: torch.Tensor,
    *,
    tol: float = 1e-4,
    distributed: bool = False,
    group: Optional[dist.ProcessGroup] = None
) -> torch.Tensor:
    """Compute the Sinkhorn balance loss, which is the mean KL divergence of the Sinkhorn balanced
    logits from the original logits.

    Args:
        L: The logits, shape (..., N, M). If `distributed` is `True`, the logits are sharded on
            dimension -2 across ranks (each rank's input is a shard).
        tol: The tolerance for the Sinkhorn iterations.
        distributed: Whether to use distributed Sinkhorn iterations.
        group: The process group to use, if `distributed` is `True`.

    Returns:
        The Sinkhorn balance loss, shape (...). If `distributed` is `True`, the true minibatch loss
        is the average of the returned losses across ranks.
    """
    world_size = dist.get_world_size(group=group) if distributed else 1
    Lq = F.log_softmax(L, dim=-1)
    with torch.no_grad():
        Lp = log_sinkhorn(Lq, tol=tol, distributed=distributed, group=group)
        Lp = Lp + math.log(Lp.shape[-2] * world_size)
    return torch.mean(torch.sum(torch.exp(Lp) * (Lp - Lq), dim=-1), dim=-1)