I've been researching tensor networks as a more interpretable architecture, but whenever I tell people this, they always ask "But is it any good?"

So I trained multiple 500M parameter LLMs on fineweb, showing the tensor variant needed ~4% more batches of data to match CE-loss.

There's a few caveats, so my personal estimate is around 15% worst to 10% better. Details below.

The Architecture

Replacing MLP w/ a Bilinear Layer

An MLP is a linear encoder, ReLU, then linear decoder.

M L P (x) = D (R e L U (E (x)))

A bilinear layer asks "what's better than one encoder? Two!"

B i l i n e a r (x) = D (L x ⊙ R x)

Where $⊙$ means "element-wise multiply" eg

[1, 2, 3] ⊙ [1, 2, 3] = [1, 4, 9]

A SwiGLU Layer (Swish Gated Linear Unit) says "Let's add in nonlinearities"

S w i G L U (x) = D (s w i s h (L x) ⊙ R x)

SwiGLU is a SOTA architecture & Bilinear is a tensor network.

Replacing Softmax Attn w/ Bilinear Attn

For a tensor network, we are only allowed polynomial nonlinearities. For attention, this means we need to replace softmax w/ a polynomial. The simplest is an element-wise squaring of the attention pattern with itself.

A t t n^{2} = (x Q K^{⊤} x^{⊤})^{2}

= (x Q K^{⊤} x^{⊤}) ⊙ (x Q K^{⊤} x^{⊤})

Notice how this is like computing the same attention pattern twice, then element-wise multiplying.

We can instead be slightly more expressive by training two sets of Q & K:

B i l i n e a r_A t t n = (x Q_{1} K_{1}^{⊤} x^{⊤}) ⊙ (x Q_{2} K_{2}^{⊤} x^{⊤})

Experiment & Results

I forked an older commit of modded-nanoGPT and trained four ~500M parameter LLMs (~GPT-medium size) on fineweb, switching out Bilinear & SwiGLU for the MLP-component and Softmax Attn & Bilinear attn for Attention. I trained on CE loss, comparing when they reached a loss of 3.0.

Here, using the Bilinear or SwiGLU layer were basically the same, but switching to bilinear attn came at cost (but a very small one. Though note the 10% more parameters for Bi_attn)

I initially ran experiments in the GPT-2 small range:

Which were much worse for the tensor-variants.

One might think "Oh, the tensor-variants actually scale better?", but I think it's because I forked from modded-nanogpt who's hyperparams & architecture is overfitted to the GPT-2 size model for Softmax attention. You'd need to do a larger hyperparam sweep (& various architecture changes, yikes) to get a better idea!

Caveats:

(1) Softmax attention ran faster cause it has a CUDA kernel

Softmax attention ran much faster because I was using F.scaled_dot_product_...

cat-state

cat-state