[ML] Point Cloud Transformer

Terminology

modulate
- In ML, modulation means changing one signal using another signal.
- So Attention $output=\alpha \odot v$ with elementwise product is modulation on value features. The value is scaled first before being summed over
MLP (Multi-Layer-Perceptron)
- A feedforward network made of linear layers + non-linearities. Technically, we are using Neurons, not just perceptrons (Rosenblatt, 1958)
- So Conv1d + GroupNorm + ReLU + Conv1d is a valid MLP

Point Transformer vs Traditional Transformer

Overview

In a point cloud, our input is the raw point cloud. Assuming we have B batches, M points in each batch. P is [B, 3, M]. Each point is

\[p_i = [x_i, y_i, z_i]\]

we first calculate K Nearest Neighbors (KNN) for each point. N is the raw neighbor point set and is [B, 3, M, K]

\[n_i = knn(p_i)\]

Then we apply self attention to each point and their K neighbors. This is to learn the structural features like edges and corners of the point cloud. Each input point $p_i$ is first mapped to an embedding $\mathbf{x}_i \in \mathbb{R}^d$ . The total x tensor is [B, D, M]

\[x_i = MLP(p_i)\]

Then, we feed $x_i$ as the only input to the self-attention unit. linear projections generate query, key, and value features. $x_i$ embedding gets fed into the point transformer network. So $q_i$, $k_i$, and $v_i$ are learned intermediate feature vectors of a point. We normally choose q, k, v to have the same dimension, so the total $q, k, v$ tensors are: [B, H, M]

\[\mathbf{q}_i = W_q \mathbf{x}_i, \quad \mathbf{k}_i = W_k \mathbf{x}_i, \quad \mathbf{v}_i = W_v \mathbf{x}_i.\]

Now, broadcast each point $p_i$ and its feature vector $q_i$ k times to a new dimension. For the sake of readability, we stick to the notation $p_i$ and $q_i$, so the total p is [B, 3, K, M]. $q$ tensor is: [B, H, K, M]

\[\tilde{p}_i = \begin{bmatrix} p_i & p_i & \cdots & p_i \end{bmatrix}_{3 \times K} \quad \tilde{q}_i = \begin{bmatrix} q_i & q_i & \cdots & q_i \end{bmatrix}_{H \times K}\]

Calculate positional encoding $\delta_{i,j}$ between point coordinates $p_i$ and its K neighbors, so distance and direction information is encoded. $\delta$ is [B, H, M, K]:

\[\delta_{i,j} = MLP(p_i - n_{i,j})\]

For each point $i$, attention is computed over its KNN neighborhood $\mathcal{j}$ (so it’s not global). we can calculuate a logit which is a vector. Logit is [B, H, M, K]:

\[l_{i,j} = MLP(\mathbf{q}_i - \mathbf{k}_j + \delta_{ij})\]

So, for each channel, we can calculate a softmax probability where $\delta_{ij}$ encodes relative positional information and a small MLP creates an additional learnable mapping for the logit to increase expressiveness and stablizes training.

\[\alpha_{ij} = \mathrm{Softmax}_j \big( MLP(\mathbf{q}_i - \mathbf{k}_j + \delta_{ij}) \big),\]

Then, the output of the attention is vector attention that uses elementwise product $a \odot b$, which modulates features in $v_j + \delta_{i,j}$ per channel. Output is [B, H, M]:

\[\mathbf{y}_i' = \sum_{j \in \mathcal{N}(i)} \alpha_{ij} \odot \big( \mathbf{v}_j + \delta_{ij} \big),\]

Discussions

What Does a Channel Represent?

Think of each channel as detecting something different:

Channel 3 → curvature
Channel 7 → edge direction
Channel 12 → density pattern

Now different neighbors may contribute differently per channel. So now, channel 3 might prefer the left neighbor. channel 7 might prefer the above neighbor. channel 12 might prefer the closest neighbor.

Of course, in a real network, each channel could be something more extract.

Positional Encoding $\delta_{i,j}$ is relative position

Vanilla transformer:

x = token embedding + absolute positional encoding

Then you do

q = Wq x, k = Wk x, v = Wv x

In a Point Transformer, position embedding is relative:

\[\delta_{i,j} = h(p_i - p_j) \\ \phi(q_i, k_j) = \gamma (q_i - k_j + \delta_{i,j})\]

Vector Attention & Elementwise Scalar Product

Recall that a regular attention logit and its probability are a scalar. Given query $q_i \in \mathbb{R}^d$, key $k_j \in \mathbb{R}^d$, and value $v_j \in \mathbb{R}^d$:

\[\phi_{i,j} = q_i^\top k_j\] \[\alpha_{i,j} = \operatorname{softmax}_j\!\left( \frac{\phi_{i,j}}{\sqrt{d}} \right) = \frac{\exp\!\left( \frac{q_i^\top k_j}{\sqrt{d}} \right)} {\sum\limits_{l} \exp\!\left( \frac{q_i^\top k_l}{\sqrt{d}} \right)}\]

Point Transformer (Zhao et al., 2021 ICCV) modified this. Instead of using dot product, the vector attention uses hardamard product

The output is a $C$-dimensional feature vector. For each channel $c$, the aggregated feature at point $i$ is computed as a weighted sum over its k neighbors:

$y_i^{c} = \sum_{k \in \mathcal{N}(i)} \alpha_{ik}^{c} \, v_k^{c},$ where $\alpha_{ij}^{c}$ denotes the attention weight assigned to neighbor $j$ for channel $c$, and $v_j^{c}$ is the $c$-th channel of the value feature at point $j$.

Main benefit for using vector attention is per-channel (where a channel is a feature dimension) weighted sum of attention weight and feature vector.

Example:

We have 2 channels and 2 neighbors.

\[V = \begin{bmatrix} 2 & 1 \\ 4 & 2 \end{bmatrix} \in \mathbb{R}^{2 \times 2}\]

Rows = channels $C$
Columns = neighbors $K$

\[\mathbf{v}_1 = \begin{bmatrix} 2 \\ 4 \end{bmatrix}, \quad \mathbf{v}_2 = \begin{bmatrix} 1 \\ 2 \end{bmatrix}.\]

Attention Tensor (Vector Attention): each channel has its own attention distribution over neighbors: $\boldsymbol{\alpha} = \begin{bmatrix} 0.25 & 0.75 \\ 0.90 & 0.10 \end{bmatrix} \in \mathbb{R}^{2 \times 2}$

Row c = attention weights for channel c
Column k = weight for neighbor k

where row $c$ contains the attention weights for channel $c$, and column $k$ corresponds to neighbor $k$. Vector attention performs an elementwise (Hadamard) product per channel and sums over neighbors:

$\mathbf{y} = \sum_{k=1}^{K} \boldsymbol{\alpha}_{:,k} \odot \mathbf{v}_{:,k}.$ Equivalently, channel-wise: $y_c = \sum_{k=1}^{K} \alpha_{c k} \, v_{c k}.$

Channel 0:

\[y_0 = 0.25 \cdot 2 + 0.75 \cdot 4 = 3.5\]

Channel 1:

$y_1 = 0.90 \cdot 10 + 0.10 \cdot 20 = 11$
Final output vector: $\mathbf{y} = \begin{bmatrix} 3.5 \\ 11 \end{bmatrix}$

Vector Attention