Computing ANDs in bilinear layers

Bilinear layers

Bilinear layers are an alternative to the standard MLP layer with some nonlinear activation function. For an input

$x$, a bilinear MLP would be

$$\begin{array}{rl}ML{P}_{bilinear}(x)& =(Wx)\odot (W^{\prime }x)\\ & =\sum _{jk}{W}_{ij}{W}_{ik}^{\prime }{x}_j{x}_k\end{array}$$ where

$\odot$ represents element-wise multiplication and

$W,W^{\prime }$ are weight matrices. Unlike other nonlinear layers, the bilinear layer can be written in terms of only linear operations and a third order tensor,

$W_{ij}{W}_{ik}^{\prime }$.

Pairwise Feature Interactions

The advantage of the bilinear layer is that it's the simplest type of nonlinearity, only having pairwise interactions. This can make it simpler to understand how input features might interact to create output features.

For example, if the input is composed of feature vectors,

$f$, with (sparse) activations,

$a$, such that

$x_i=\sum _{\alpha }{a}_{\alpha }{f}_{\alpha i}$, then the output can be written as

$$\begin{array}{rl}ML{P}_{bilinear}(x)& =\sum _{\alpha \beta }{a}_{\alpha }{a}_{\beta }\left(\sum _{jk}{W}_{ij}{W}_{ik}^{\prime }{f}_{\alpha j}{f}_{\beta k}\right)\end{array}$$ after grouping activations and features together.

If we expect output features to represent pairwise operations of input features (eg ANDs of input features), then it's natural to define the output activation and features as:

$$\begin{array}{rl}{a}_{(\alpha \beta )}^O& \equiv a_{\alpha }{a}_{\beta }\\ f_{(\alpha \beta )i}^O& \equiv \sum _{jk}{W}_{ij}{W}_{ik}^{\prime }(f_{\alpha j}{f}_{\beta k}+f_{\beta j}{f}_{\alpha k})\\ ML{P}_{bilinear}(x)& =\sum _{(\alpha \beta )}{a}_{(\alpha \beta )}^O{f}_{(\alpha \beta )i}^O\end{array}$$ where

$(\alpha \beta )$ indexes over pairs of input features, including twin pairs like (00). The features are written in a form that is symmetric over

$\alpha ,\beta$.

Computing ANDs

If the input activations

$a_{\alpha }$ are sparse and boolean, then

$a_{(\alpha \beta )}^O\equiv a_{\alpha }{a}_{\beta }$ already computes the AND for the pair of input features. The feature vector

$f_{(\alpha \beta )}^O$ then represents the AND in the output space. The original input features are also represented by

$f_{(\alpha \alpha )}^O$.

The main constraint is that the output features form an over-basis, meaning that the dot product of any two features is sufficiently small.

It's possible that random feature vectors (coming from random

$W$ and

$W^{\prime }$) are enough for an overbasis. For a

$d$-dim space, the dot product of two (normalized) gaussian random vectors is a random value with mean zero and standard deviation

$\approx \frac{1}{\sqrt{d}}$.

To read-off the value of the AND of two features, we could simply dot product the MLP output with the feature vector for the AND we want to extract and normalize. This means dotting with

$f_{(\alpha \beta )}^O/||f_{(\alpha \beta )}^O|{|}^2$. For random gaussian feature vectors, this would give

$a_{(\alpha \beta )}^O$ (equal to

$a_{\alpha }\wedge a_{\beta }$) plus inteference terms from the other active feature vectors. For

$\ell$ input features active, the interference terms will be of order

$\frac{\ell }{\sqrt{d}}$ after adding up

$\pm \frac{1}{\sqrt{d}}$ over

${\ell }^2$ active feature pears. For

$\ell =5$ and

$d=1000$, the interference terms are on the order of

$0.16$. Reading off XORs can be done similarly, using the fact that

$a_{\alpha }\text{XOR}{a}_{\beta }=a_{\alpha }+a_{\beta }-2a_{\alpha }\wedge a_{\beta }$.

Random

$W$,

$W^{\prime }$ might be enough to represent all pairwise feature ANDs. The limitation is that no two random vectors should have high interference. The appendix shows that the probability of two random vectors having a cosine similarity of more than

$\delta$ is

$$P[\frac{|x^{T}y|}{||x||||y||}\ge \delta ]\approx \sqrt{\frac{2}{d{\delta }^2\pi }}\exp (-d\frac{{\delta }^2}{2})$$ when

$\delta \gg 1/\sqrt{d}$. For example,

$d=1000$ and

$\delta =0.3$ gives

$P\sim {10}^{-21}$.

Overall the number of ANDs that could be stored before two feature vectors become similar is exponential in

$d$. Inteference during read-offs is small as long as the number of active features,

$\ell$, is small compared to

$d$.

In practice, features (and ANDs of features) will have correlations, so a model could learn to minimize the interference between features that are often active together. If feature activity is power law distributed, there is potential to greatly reduce inteference by carefully choosing the feature vectors of the most active features.

Non-boolean features

Even if the activities

$a_{\alpha }$ are sparse but not boolean, our construction of output features still appears useful. The output activations would be

$a_{(\alpha \beta )}^O=a_{\alpha }{a}_{\beta }$ which can serve as a basis for building up some nonlinear function of the initial features.

The feature vectors that the model learns may also play a role. For example if three pairs of features all have similar output feature vectors, say

$f^O$, then the model may be effectively computing

$a_1{a}_2+a_3{a}_4+a_5{a}_6$ and representing it with

$f^O$.

Appendix

Probability of high interference: According to these lecture notes on high-dim data [link] (chapter 4), the dot product of a gaussian random vector

$x$ (entries iid with

${\sigma }^2=1$) with a normalized random vector

$\frac{y}{||y||}$ is also gaussian distributed with variance 1. Therefore

$$P[\frac{|x^{T}y|}{||y||}\ge \epsilon ]=\frac{2}{\sqrt{2\pi }}{\int }_{\epsilon }^{\mathrm{\infty }}{e}^{-\frac{t^2}{2}}dt\equiv \text{erfc}\left(\frac{\epsilon }{\sqrt{2}}\right).$$

The norm of

$x$ concentrates around

$\sqrt{d}$ with small fractional deviations of order

$1/\sqrt{d}$. Neglecting these deviations, the dot product of two normalized gaussian vectors is then

$$P[\frac{|x^{T}y|}{||x||||y||}\ge \delta ]\approx \text{erfc}\left(\delta \sqrt{\frac{d}{2}}\right)$$ using the equation above with

$\delta \equiv \frac{\epsilon }{\sqrt{d}}$. For

$\delta \gg 1/\sqrt{d}$, this gives

$$P\approx \sqrt{\frac{2}{d{\delta }^2\pi }}\exp (-d\frac{{\delta }^2}{2}).$$

For

$d=1000$, a high interference of

$\delta =0.3$ has probability

${10}^{-21}$ and for

$\delta =0.5$ the probability is

${10}^{-56}$.