PointNet pooling: one summary for a set

A classifier needs one fixed-size vector, but a cloud has any size and no order

Say a network must name the object a point cloud shows: chair, lamp or plane. Its last layers expect one input vector of a fixed size. A cloud, though, can hold 500 points or 50,000, listed in any order. Writing all the coordinates into one long vector fails twice: the size changes from cloud to cloud, and the same shape listed in another order gives a different vector. As the permutation invariance unit explains, the summary must be permutation invariant.

PointNet, a network published in 2017 by Qi and colleagues, solves this in two steps.

Same MLP for every point, then the maximum of each feature

1. Shared MLP. One MLP h, with one set of weights, maps each point p on its own to K features. A cloud of n points becomes a table of n rows and K columns. The rows follow the points: shuffle the points and the rows shuffle with them.
2. Max pooling. For each column k, keep the largest value: g_k = max over all points p of h_k(p). The table becomes one vector g of K numbers: the summary of the cloud.

The maximum of a list does not depend on the list's order, and the shared MLP treats every point the same way. So g is the same for every order of the points. It holds K numbers, whatever n is. A second MLP then reads g and names the object.

Worked by hand: four points, three features

Use 2D points to keep the numbers small: A = (0, 0), B = (2, 0), C = (1, 2) and D = (1, 0.5). Let the shared "MLP" be one linear layer followed by ReLU, which keeps positive numbers and turns negative ones into 0. Its three features are

h(x, y) = ( ReLU(x), ReLU(y), ReLU(y − x) )

The third feature says how far a point sits above the diagonal line y = x.

Predict first. Which point will set the largest value of each feature? Will any point not matter at all?

Shuffling the points cannot change the summary; moving a winner can

The maximum keeps the summary steady when points go missing

Any order-free way of combining the rows would work: the sum, the mean or the maximum. The PointNet authors compared max pooling with average pooling and with an attention-weighted sum on a shape classification task, and max pooling won by a large margin. They also showed why it is robust. With K features, at most K points can hold a maximum. They called these the critical points. Deleting any other point, or moving it a little, leaves the summary unchanged, so a cloud with a few missing or noisy points gets nearly the same summary.

Real sizes are larger than the example. The original PointNet used a shared MLP with layer widths 64, 64, 64, 128 and 1,024, so each cloud became a summary of 1,024 numbers. An MLP with widths 512 and 256 then picked the object's class.

pointnet_pool.py

import numpy as np

rng = np.random.default_rng(0)
W1 = rng.normal(size=(3, 64))        # a shared MLP: 3 -> 64 -> 128
W2 = rng.normal(size=(64, 128))

def h(P):                            # P holds one point per row
    return np.maximum(0, P @ W1) @ W2

def pointnet_pool(P):
    return h(P).max(axis=0)          # the largest value of each feature

P = rng.normal(size=(500, 3))        # 500 points in 3D
g = pointnet_pool(P)
print(g.shape)
# Shuffle the points: np.allclose says whether the summary stays the same, up to rounding.
print(np.allclose(g, pointnet_pool(P[rng.permutation(500)])))
print(pointnet_pool(P[:50]).shape)   # 50 points: same size
# argmax(axis=0) finds, for each feature, the point that holds its maximum.
# np.unique keeps each of those points once, and len counts them.
print(len(np.unique(h(P).argmax(axis=0))))  # how many points set a maximum

(128,)
True
(128,)
38

The last line counts the critical points: only 38 of the 500 points decide the 128 numbers of the summary.