Pairwise Distances with Broadcasting

The problem

Models represent words, documents and images as vectors: lists of $D$ numbers, where similar things get vectors that are close to each other. A very common job is to compare every vector of one set with every vector of another, for example to find which stored documents are closest to a question (retrieval, section 15). The result is a table of distances: one row per vector of the first set, one column per vector of the second.

You get two tables of vectors:

• a, with $N$ rows: row $i$ is the vector $a_i$;
• b, with $M$ rows: row $j$ is the vector $b_j$;
• both have $D$ columns, one per coordinate.

Return the $N \times M$ table whose entry $(i, j)$ is the squared distance between $a_i$ and $b_j$.

Squared distance

To compare two vectors: subtract them coordinate by coordinate, square each difference, and add up.

where:

• $a_{id}$ is coordinate $d$ of vector $a_i$, and $b_{jd}$ is coordinate $d$ of vector $b_j$;
• $\lVert v \rVert^2$ is the squared length of a vector $v$: the sum of its coordinates squared;
• it is the usual distance of geometry, without the square root (which does not change which vectors are closest).

Example. $D = 2$. a has the rows $a_0 = (0, 0)$ and $a_1 = (1, 1)$; b has the rows $b_0 = (3, 4)$ and $b_1 = (0, 1)$.

• $\mathrm{D}_{00}$: $a_0 - b_0 = (-3, -4)$, and $9 + 16 = 25$.
• $\mathrm{D}_{01}$: $a_0 - b_1 = (0, -1)$, and $0 + 1 = 1$.
• $\mathrm{D}_{10}$: $a_1 - b_0 = (-2, -3)$, and $4 + 9 = 13$.
• $\mathrm{D}_{11}$: $a_1 - b_1 = (1, 0)$, and $1 + 0 = 1$.

The result is the $2 \times 2$ table [[25, 1], [13, 1]].

Method 1: all the differences at once (too much memory)

Broadcasting can build every difference $a_i - b_j$ in one expression: a[:, None, :] - b[None, :, :] has shape $(N, M, D)$, one difference vector for every pair. Squaring and summing over the last axis gives the answer.

This is a good way to check your answer on small inputs, but it stores $N \times M \times D$ numbers. For $N = M = 300$ and $D = 128$ that is 11.5 million numbers, 92 MB (megabytes), only to keep 90,000 of them; with real data it grows to gigabytes.

Method 2: expand the square (little memory)

The square of a difference can be expanded, just like $(x - y)^2 = x^2 + y^2 - 2xy$ for numbers:

where $a_i \cdot b_j = \sum_d a_{id}\, b_{jd}$ is the dot product: multiply coordinate by coordinate and add up. Each of the three parts is cheap:

• $\lVert a_i \rVert^2$ for every $i$: one number per row of a, a column of $N$ numbers;
• $\lVert b_j \rVert^2$ for every $j$: one number per row of b, a row of $M$ numbers;
• $a_i \cdot b_j$ for every pair: exactly the matrix product a @ b.T, an $N \times M$ table.

Example. Same vectors as above.

• Squared lengths of a: $\lVert a_0 \rVert^2 = 0$, $\lVert a_1 \rVert^2 = 1 + 1 = 2$.
• Squared lengths of b: $\lVert b_0 \rVert^2 = 9 + 16 = 25$, $\lVert b_1 \rVert^2 = 0 + 1 = 1$.
• Dot products a @ b.T: [[0, 0], [7, 1]], for example $a_1 \cdot b_0 = 1 \cdot 3 + 1 \cdot 4 = 7$.

Entry $(1, 0)$: $2 + 25 - 2 \cdot 7 = 13$, the same as method 1. The whole table is again [[25, 1], [13, 1]], and nothing bigger than $N \times M$ was ever stored.

What to implement

pairwise_sq_dists(a, b): a has shape $(N, D)$, b has shape $(M, D)$. Return the $(N, M)$ table of squared distances.

Requirements

• No Python loops.
• Every distance is $\geq 0$, and the distance of a vector to itself is $0$ (up to rounding).
• Memory: one test uses $N = M = 300$, $D = 128$ and fails if memory use goes above 20 MB, so method 1 does not pass. Use method 2.

Hints

1. Write method 1 first and keep it to check method 2 on small inputs.
2. In method 2, the squared lengths of a must be a column, shape $(N, 1)$, and those of b a row, shape $(1, M)$. Added together, broadcasting turns them into the $(N, M)$ table of $\lVert a_i \rVert^2 + \lVert b_j \rVert^2$ (see Broadcasting, example 2, on the theory page).
3. Rounding can make the result slightly negative, such as $-10^{-15}$, when $a_i$ and $b_j$ are almost equal. Replace negative values with 0: np.maximum(d, 0.0).