Sliced Optimal Transport

Imagine wanting to combinatorially reassign the pixels of one image to match another image subject so some cost metric. An exhaustive solution like the Hungarian Algorithm has complexity, making it both relatively pointless as a style transfer algorithm and very expensive.

Sliced optimal transport is an algorithm for approximately solving the assignment problem in high dimensions. Given two point sets, it finds a correspondence that minimizes total displacement by reducing the problem to a series of one-dimensional sorts.

How It Works

The animation below demonstrates sliced optimal transport in two dimensions. We start with source points (blue) and target points (orange). Each iteration chooses a random direction and projects both point sets onto that line. After sorting both projections, we pair them by rank and nudge each source point along the projection direction based on the difference between its projected value and that of its paired target.

In practice, we would average updates across a batch of random directions before applying them. For illustration, we show a fractional update from a single direction per step. Repeating this process with new random directions, the source points gradually rearrange to match the target distribution.

References

Alan Wolfe provides an excellent explanation of sliced optimal transport in his article, Interpolating Color Image Histograms Using Sliced Optimal Transport. Theo Honohan pointed me to Color Style Transfer Techniques using Hue, Lightness and Saturation Histogram Matching by László Neumann and Attila Neumann for alternative approaches to color transfer.

Demo

The interactive demo below applies sliced optimal transport to rearrange the pixels of one image to match the color distribution of another.

My original code from this notebook is presented below. However, this is not the code being run in this notebook. This notebook uses a Web Worker pool to perform the batches in parallel since the algorithm can get a bit slow. You can read the modified implementation in the Notebook Kit source for this page.

function* slicedOptimalTransport(
  src,
  tgt,
  { maxIterations = 100, batchSize = 4, tolerance = 1 } = {}
) {
  if (src.length !== tgt.length)
throw new Error("Source size must equal target size");

  // TODO: I could probably avoid this if I better understood generator functions
  mutable history = [];

  const N = src.length >> 2;
  const index = new Uint32Array(N);
  const srcProjection = new Float32Array(N);
  const tgtProjection = new Float32Array(N);
  const adjustment = new Float32Array(N * 3);

  let delta = Infinity;
  let iteration = 0;
  while (delta > tolerance && ++iteration <= maxIterations) {
adjustment.fill(0);
for (let batchIndex = 0; batchIndex < batchSize; batchIndex++) {
  let [v0, v1, v2] = vec3normalize([], [randn(), randn(), randn()]);

  for (let i = 0, i4 = 0; i < N; i++, i4 += 4) {
    index[i] = i;
    srcProjection[i] = v0 * src[i4] + v1 * src[i4 + 1] + v2 * src[i4 + 2];
    tgtProjection[i] = v0 * tgt[i4] + v1 * tgt[i4 + 1] + v2 * tgt[i4 + 2];
  }

  // Native sort works for target since all we need is the sorted vector, but
  // custom sort is required for the source since we need to sort the indices
  // in parallel with the projection.
  sort(srcProjection, index, 0, N - 1);
  tgtProjection.sort();

  for (let j = 0; j < N; j++) {
    const projectedDiff = tgtProjection[j] - srcProjection[j];
    const i3 = index[j] * 3;
    adjustment[i3 + 0] += v0 * projectedDiff;
    adjustment[i3 + 1] += v1 * projectedDiff;
    adjustment[i3 + 2] += v2 * projectedDiff;
  }
}

delta = 0;
for (let i3 = 0, i4 = 0; i4 < N * 4; i3 += 3, i4 += 4) {
  const dr = adjustment[i3] / batchSize;
  const dg = adjustment[i3 + 1] / batchSize;
  const db = adjustment[i3 + 2] / batchSize;
  src[i4] += dr;
  src[i4 + 1] += dg;
  src[i4 + 2] += db;
  delta += dr * dr + dg * dg + db * db;
}
delta = Math.sqrt(delta / N);
mutable history = mutable history.concat([{ iteration, delta }]);

yield { delta, src };
  }

  return { delta, src };
}