If you're interpreting <1B parameter models, you should use a tensor transformer

To all my fellow researchers doing SLT, computational mechanics, one of ARC's programs, natural abstractions/condensation, proofs on NNs (or any interp on small models), this is for you.

Tensor transformers (ie replacing your MLPs & attention with bilinear variants) are performant and allow you to deploy the full power of linear algebra. In fact, our recent paper used generalized cosine similarity on the full tensor transformer. And yes, I mean cos-sim defined on the eg 9th order tensor, not individual vectors or matrices. This removed all the symmetries/invariances that weren't functionally relevant.

But tensor-variants don't generalize to "real models", right?

The architectures are very similar:

SwiGLU(x) = D(swish(Lx) ⊙ Rx) (used by DeepSeek-V3, Kimi K2, and Qwen3))
Bilinear(x) = D(Lx ⊙ Rx) (this is the tensor version)

Where D, L, & R are linear matrices. For reference:

MLP(x) = D(ReLU(Lx))

Due to the double-encoder/multilinearity, SwiGLU & Bilinear have no global Lipschitz constant (and other similar inductive biases). This means results like finetuning away the normalization might not generalize to these SOTA archs since this was only run on single-encoder MLPs.

For attn, the more SOTA tensor-arch is:

Bilinear_Attn = OV()

Compared to softmax attention, this does produce denser attention patterns (with softmax producing sparser). This ends up being a better inductive bias for eg chess or othello, so I do expect more differences here than in the bilinear layer.

You can still use RMSNorm (which is secretly a tensor network) and residual streams (which are secretly a tensor network). You can even include Mixture of Experts (which AREN'T tensor networks, but they only exponentially blow up the number of possible paths a little bit. But as a wise man once said "Compositionality is a spectrum", so you can still get most of the benefits of both).

Frontier Models aren't the Only Thing That Matters

I don't think getting frontier models to be tensor networks is the main value add. Current frontier mdoels aren't robust and can't shouldn't be deployed in high stakes settings (kind of our whole problem if you think about it).

If we can reverse engineer models, we can solve deep learning. We can get the data and clarity to know how data + arch --> algorithms and have more control on how the model generalizes.

We can create task-AI that can be robustly deployed, allowing safe, continued economic growth. If we have an international pause on AI, we can still share the fruits of task-ai without sharing algorithmic secrets.

We can interpret bio-models, discovering more accurate gene regularity networks to develop more targeted medicine.

And if we have safer AI that's just more expensive to train, then it'd be good to have that shown clearly as we continue to get warning shots! Right now I can't even interpret a GPT-2 small sized tensor transformer, so I can't present them as a safe alternative yet. I believe we should focus on solving that first.

My Extreme Pessimism (or Ignorance)

I'm trying to reverse engineer a language task (and some toy algorithmic tasks) on tensor networks AND IT'S STILL REALLY HARD. I have all these advantages and am still having a difficult time. Bilinear layers/ tensor transformers are the easier case, and if you're not able to solve your task there, then you're fundamentally confused. BUT fundamentally confused about an easy-to-analyze object.

If you'd like to chat about how tensor transformers can help with your research agenda (only safety-related ones please), do book a call or dm me on discord at # loganriggs.

  1. Conservatively 90% as efficient as a normal transformer
  2. We never instantiate the eg 9th order tensor and there are efficient algorithms to compute this.
  3. Due to the element-wise multiplication, GLU architectures have an inductive bias towards anti-podal representations. In plain English: if we look at one decoder row, this is scaled by some scalar depending on the encoder.So we're just scaling the vector d by a scalar. If the two intermediate scalars, & , are the same signs, then it scales positively. If they're opposite signs, then they scale negatively. So d can store two things that are mutually exclusive to each other (ie in an antipodal fashion), which might be a decent inductive bias for language. All GLU architectures, SwiGLU & Bilinear layer included, have this property.
  4. There's also other tensor attention variants we could invent, I just think it's not the highest priority to focus on now.
  5. AI that just does 1 thing. You can think of SGD as a program search, then with interp as extracting the program that does eg electric grid, banking, farming, etc
添加评论
点赞收藏
点踩分享查看原文
评论
?
参与讨论