Self Attention and Pooling in the Set Transformer
I have been interested in machine learning (ML) for a long time (recently dusted off my Kaggle account to figure out it is 9 years old!). It started from taking one of the first ever online courses on Coursera by Andrew Ng back then. Since that time I have always been curious how ML and deep learning models can process sets of data with varying number of observations and how the models can be invariant to the order of the observations in the set.
Certainly, there are ways to overcome the problem of sets. The straightfoard approach is to use some form of pooling, such as averaging, summing, max pooling. These operations can be either done on the initial observation sets or somewhere inside the model after e.g., few hidden layers. The drawback of these pooling methods is partial loss of data such as interdependincies between observations in a set.
Set Transformer\cite{cite1-6} is one recent architecture that was designed to work with sets. And I decided to spend few ours on a weekend to find out how Set Transformer can deal with the varying number of observation within each set. The Set Transformer employs attention blocks to work and this post points out the important idea that there are two distinct types of the attention blocks although the structure of these blocks is the same.
The first type of attention block is the self attention within the set. The self attention has nothing to do with the "smart" pooling in the first place. The self attention blocks (SAB) aim into learning the interdependincies within the set observations either using the multihead attention blocks (MAB) or the induced SABs (iSAB,to decrease the dimensionaly of the MABs). These SABs and induced SABs apply the attention on the observation within the set and learn correlations between the observations. One dimention of SAB or iSAB still equalt to the n number of observations in the set.
What enables the smart pooling and removal of the n dimentionality in the set transformer is the second type of the attention block called pooling by the multihead attention (PMA). PMA uses the same MAB as in SAB or iSAB, but instead of learning the correlations between observations, it takes the output of those learned correlation to learn how to pool them in a smart way to remove the n dimentionality. In the paper\cite{cite1-6}, PMA works in the following way:
\begin{equation}\text{PMA}_k(Z) = \text{MAB}(S, \, \text{rFF}(Z)),\end{equation}
where \(S \in \mathbb{R}^{k \times d}\) is the learnable seed matrix (with \(k\) seed vectors and \(d\) is the dimentionality of the one observation vector), \(Z\in \mathbb{R}^{n\times d}\) is the output of the SABs or iSABs that learned correlations between set observations, and \(\text{rFF}\) is the row-wise feed-forward layer.
To see how the \(n\) dimenationalty disappears, let us use the definition of the multihead attention block\cite{cite2-6}
\begin{equation}\text{MAB}(S,\,\text{rFF}(Z)) = \text{Multihead}(S,\,\text{rFF}(Z),\,\text{rFF}(Z); \,\omega) = \omega(S(\text{rFF}(Z))^\text{T})\text{rFF}(Z),\end{equation}
where \(\omega\) is the activation function.
Consider for simplicity that \(\text{rFF}(Z)\) does not change the dimentions of \(Z:\, Z\in \mathbb{R}^{n\times d}\) and \(\text{rFF}(Z)\in \mathbb{R}^{n\times d} \). The dimentions \(\text{MAB}(S,\,\text{rFF}(Z))\) of are then \(\mathbb{R}^{k \times d} \cdot \mathbb{R}^{d \times n} \cdot \mathbb{R}^{n \times d} = \mathbb{R}^{k \times d}\). The \(n\) dimentionality disappeared during the matrix multiplication.