Solving the "Very Selective Copying" problem with a Very Small Language Model ←

This is a very numeric continuation of a previous experiment. To get a grip on the details, please check "Selective Copying" first.

Basically, it's a synthetic question-answer dataset that requires some "computational" skill. See, if you can find out the rule by yourself:

NA: 2>1: AN   
DQ: -1: Q    
EJWHB: 3>5: EJBHW
ULHP: 3>4, 2>1: LUPH 
YNP: 3>1, -3, 2>1: NP   
EJACQ: 1>2, -1, +3: EAJCQ
YESBR: 3>5, 5>1, 1>5, 3>4: YEBRS
UXMP: -2, -3, 1>2, +2: MPU  

Spoiler: First few letters are the program input, the numbers and operators are the program and answer, after the second colon, is the result of the program.

• x>y means: exchange xth item with yth item.
• -x means: remove xth item from memory and put it onto the stack.
• +x means: pop last item from stack and insert at position x

I argue, that this dataset

• requires a large enough receptive field
• requires some cognitional ability to build up a history and to use it subsequently.
• should actually be called Selective Swapping

Preliminary tests ←

Just to get a feeling for the dataset, i took the best model from the "Selective Copying" experiment with a few parameter variations and ran it on questions that have an input length of 2 to 5 items and 1 to 5 operations:

Training now takes 4 million steps, and could even be a bit longer. But generally, the loss curves seem to converge at about there.

I'm using dilation all the time, because 1.) it's required for the receptive field of the 6-layer network and 2.) because it runs faster. For the 12 and 18-layer networks i just expanded the dilation settings without much thinking about it. (Actually i did think quite a bit about it but decided to evaluate good 12 or 18-layer dilation settings another time)

Here are examples from the validation set from the worst and the best performing network:

6-layer/32-chan: validation mask error 17.6% ←

It is noticeable that long programs seem to be the problem for this network.

12-layer/128-chan: validation mask error 0.7% ←

Wow! Everything correct, even the long ones!

Number of operations versus number of layers ←

The above example is a bit messy because it has a variable number of operations and we don't know how many answers are fully correct, only how many characters. So first, we add a new evaluation metric validation_sample_error%. While the mask error gives the percentage of wrong characters within the mask area, the sample error gives the percentage of wrong answers, even if only one character is wrong.

Further, to better evaluate the relationship between number of operations (length of the "program") and the number of layers in the network, i set a fixed number of operations.

I also dropped the stack operations (+/-) for now and made it all a bit tighter because i'm not interested in the influence of the receptive field in this experiment. Example data with 5 operations looks like this:

NPFKT:5>4 3>1 1>4 4>2 1>3:NFTPK
UKLRM:5>1 4>2 1>4 2>4 3>2:KLMRU
LCEWI:3>5 3>2 4>2 3>1 3>5:CWEIL
LRUPX:1>4 3>2 5>4 1>3 2>1:URPXL
UQJOR:5>4 2>5 5>4 4>3 2>5:URQJO
BYGMJ:5>4 4>3 2>3 5>2 5>3:BMJGY
PBNFM:3>5 5>1 1>3 2>3 5>3:MNPFB
TNEOL:1>3 4>1 4>1 4>3 5>3:ENLTO
ALEVT:1>2 2>1 2>1 4>2 4>3:LVAET
UIZNQ:3>4 2>3 5>3 4>1 1>5:INQUZ

Running the networks on 2, 3, 4 and 5 operations per question:

6-layer ←

12-layer ←

18-layer ←

And here is it all together in a table. The dilation settings are left out for readability. They are 1,2,3,4,5,1 for the 6-layer, 1,2,3,4,5,1,2,3,4,5,1,1 for the 12-layer and 1,2,3,4,5,1,2,3,4,5,1,2,3,4,5,1,1,1 for the 18-layer network.

Something is obviously very strange. None of the networks can handle 5 operations. But some of them just did in the last experiment??

What i really like about this neural net research is that one often gets along with intuition only. We simply don't know enough until all tests are made, which probably never happens. However, in this case, i would argue that a 5 operations sequence is just too long for the network to learn the task in reasonable time. Put yourself into the eyes of the model. The training algorithm constantly tells you:

If you see this: PHMTW:5>3 2>3 5>4 5>3 1>3:????? you should output that: PHMTW:5>3 2>3 5>4 5>3 1>3:TWPMH

on and on ... for millions of examples. Now, for a human, maybe a programmer even, this should be solvable pretty quick. Once you have seen 10 examples, you can be sure what the algorithm is. But for a neural network trained with stochastic gradient descend it is basically a brute-force approach. It's like: Hey network, you gotta lot of parameters, that's fine, now if you see this X, and you should output that Y, there is a chance that inreasing this one parameter while decreasing that other one would help you to output the correct answer, the next time you see the same question. And on and on...

Once again, reasoning from intuition, the network can probably learn the task much easier when it has some easier examples in the dataset. Schools and such also don't teach children algebra by starting with: What is the answer to 7 + 4 + 9 + 2 + 6 = ??

Following, the training set contains questions with 2 to 5 operations, e.g.:

HGOKX:3>1 5>4:OGHXK
CXZMN:5>4 2>5:CMZNX
RMYTC:2>1 4>5 4>2:MCYRT
DNAHL:1>3 3>4 1>2 4>1 3>2:DHANL
HSBQA:4>5 4>5:HSBQA
DTFKS:4>5 3>2 4>5 5>1:SFTKD
UXGJN:4>2 1>2 1>4 3>1 3>2:GXUJN

and the validation set still holds the same 5-operations-only questions for which all networks have failed previously.

Here's the loss and error curves of the 18-layer network trained with 5 and 2-5 operations:

The bright one is the brighter one. Indeed, the 18-layer gets down to 2.5% sample error. It is able to solve the task (to a degree), it could just not learn it from the previous dataset. One could argue, though, that it would eventually learn the task just from the 5-operations examples but it would probably take 10 to 100 times more computation / CO2-emissions / floodings / draughts / you-name-it.

We can see from the table, that the 12 and especially the 6-layer networks are struggling. Looking at the plots of the 6-layer networks trained with 5 and 2-5 operations, we can see that the mask error decreases by a good amount but actual sample error stays roughly the same. It learned to put some letters at the right place but still fails for almost every validation sample:

• Quick takeaway: Put also some easy examples in the training set!

The curves suggest, however, that training has not yet converged. There are sure a few more per-mille to squeeze out.

This is a good point of origin for further experimentation. Can we get the 6-layer network to solve the 5-operations Very Selective Copying problem,

• without adding so many modules that it actually resembles a 12-layer network
• without making it slower to execute than the 12-layer network
• bonus: by keeping the number of model parameters equal or even lower

In other words, is there a trick, maybe to pass data around in a different way, that strongly increases the computational performance?

Quick comparison with Mamba and LSTM ←

Just for another set of baselines, i tried the state-spaces/Mamba (yet the slow version because the fast doesn't run on my system) and the all-beloved LSTM (the pytorch implementation).

The yellow/brownish is the 6-layers model from above, green is Mamba and blue the LSTM. All of them have 6 layers

• None of these especially-crafted models reached a significant performance gain.
• To be fair: The Mamba model is very small compared to the Conv1d, just 28% of parameters. Though it manages to consume 7.8x more computation time. The first tiny drop in sample error occurred after 1 hour of training and i do not have the patience today to train an equally-sized Mamba for comparison. (The dynamics of the curves suggest, that it would not change much)
• The LSTM basically archived nothing (though equal-sized). It's about as bad as the 6-layer Conv1d with only 5-operations questions in the training set. At least, it's blazinlgy fast!
• Disclaimer: I do not know any good practices about using or training LSTMs or Mambas and surely made some grave mistake..

Back to the residual CNN. To find some interesting method, quick, i switched to 3 operations and 3-layer networks. The above condition holds as well. The 3-layer network is not able to solve the 3-operations task in the validation set, but it only takes about 2 minutes to find out:

The validation sample error stays roughly at 100%.

I tried a couple of things, like adding cross-residual connections, e.g. from layer 1 to layer 3, which did not noticeably improve performance. Also tried learnable weights, so that each layer can use as input a weighted mix of all previous layers. Was not much help either, also not for the 6-layer networks.

Attention, please! ←

UPDATE: After revision of the used experiment files, i just painfully realized that some of the below findings are wrong. In short, i forgot to put the QKV and the attention activation variables to the network constructor and as such, the tables comparing these values show something completely different. There is obviously so much deviation between similar runs, that the results fooled me to make conclusions about best parameters, although there was nothing to compare...

I'm leaving the manuscript below as it is, as a reminder how stupid the search for explanations was, without a scientific basis. However, the corrections are linked at the appropriate places.

Eventually i remembered that, in an age-old lecture by Geoff Hinton, he argued that for this particular case he showed, the network can not succeed without a cell-wise multiplication somewhere. Then there is this relatively recent Schlag, Irie & Schmidhuber paper (2102.11174) "Linear Transformers Are Secretly Fast Weight Programmers", which states: Fast Weight Programmers (FWPs, 1991) program their fast weight memories through sequences of outer products between self-invented key and value patterns. I just like the self-invented notion.

To make it short, i plugged a torch.nn.MultiheadAttention module into the last layer, such that the 1d-convolution, which normally outputs 64 channels, now outputs 128 channels, which are split into 64 query and 64 key channels and the attention-values are the input of the layer. The pseudo-code looks like:

qk = Conv1d(64, 128)(inp)
q, k = split(qk, 64)
outp = attention(q=q, k=k, v=inp)

Not shown is, that q, k and v (which are each of shape [B, C, L]) are transposed before the attention to move the channels (C) to the last dimension ([B, L, C]), and make it a channel-wise attention.

The pytorch MultiheadAttention has a predefined fixed length, because it adds bias and a couple of other learnable parameters, so it has to be used channel-wise, since we don't want to fix the sequence length that the network processes.

But wait, the Fast Weight Programmers-paper, said self-invented keys and values not self-invented keys and queries. So, for completeness, let's test all the combinations

self-invented keys and values ←

kv = Conv1d(64, 128)(inp)
k, v = split(kv, 64)
outp = attention(q=inp, k=k, v=v)

self-invented queries and values ←

qv = Conv1d(64, 128)(inp)
q, v = split(qv, 64)
outp = attention(q=q, k=inp, v=v)

self-invented queries, keys and values ←

qkv = Conv1d(64, 192)(inp)
q, k, v = split(qkv, 64)
outp = attention(q=q, k=k, v=v)

Unfortunately, the validation sample error varies quite a lot for each run of the same network:

So i'm doing 5 runs each and report the average. That's 80 runs for this messy plot and the table below!

Each run has the same validation set but different training sets and initial network weights. (reproduce with convtext-qa-program-3ops-attn-invent.yml @ 613064ad)

UPDATE: Please check the correction!

• qkv is one of the self-invented methods above. e.g. QK means, Query and Key come from convolution, Value is unprocessed.

• attn: 0,0,0 is no attention

• attn: 0,0,X is multi-head attention in final layer with X heads

• attn: 0,0,True is a self-built attention, loosely following equation 6 in Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention (2006.16236]). In contrast to the torch MultiheadAttention, The self-made function does not have any parameters and works on sequences and channels of any length.

  def attn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
      q = F.elu(q) + 1.
      k = F.elu(k) + 1.
      return q @ (k.permute(0, 2, 1) @ v) / (v.shape[-1] * v.shape[-2])
  

• (std) is the standard variation given for the validation sample error

So,

• the average validation sample error is smallest for the self-built, self-invented QK attention, followed by all the other self-built variants.
• Best number-of-heads for pytorch attention seems to be 8,
  • and the only performing self-invented combination seems to be Key/Value, as stated in the FWP paper..
• The absolutely-prooven-to-not-working model is pytorch attention with QK and heads=1.
• Generally, the pytorch attention module is very unstable in this experiment. Look at the error deviations! Occasionally, one of 5 runs produced a far better validation sample error than 0.8%.

Admittedly, there is something about this so-called Attention that works pretty well. Thinking in terms of the receptive field, these two dot-products in the attention formula bring every point in the sequence in contact with every other point, in a way that makes some sense. How exactly it makes sense, i can not tell. That's why i run these tests.

Compare attention activation function ←

The above activation function in the self-built attention module, elu(x) + 1, is suggested by the authors (2006.16236), mainly because it needs to be >= 0 and the elu gives a smooth gradient (unlike relu).

The Fast Weight Programmers-paper (2102.11174) suggests to replace this function with a more sparsity & orthogonality - aimed Deterministic Parameter-Free Projection (DPFP). Comparing the two activations for the 3-operators, 3-layers, 1-million-steps setup, averaged over 5 runs each, gives:

UPDATE: Check the correction!

Well, there is certainly a more powerful activation to be found for this particular task..

Compare position of self-attention in 3-layers network ←

Placing the best (most consistent) attention method from above into all the layers:

Well, how can we explain that? It's very unfortunate that there is no obvious rule popping out, where to put the attention layers, except that there should be one in the final layer.

Or do you instantly know how to setup a 5-layer network for the 5-operations task, now?

Compare position of self-attention in 5-layers network ←

For efficiency, keep the above dataset (and the 1 million steps of training), and use a 5-layer network with dilation=[2, 3, 4, 5, 1] to try all the attention-permutations.

Okay okay, there is some slight pattern for arrangement of the self-attention modules emerging.

Now pick the best ones and train and test against 5-operations questions (at 4 million training steps):

It's still not completely sufficient. Let's switch back to 6-layer networks, again copying the best attention-settings, although, there is only one performing okay-ish, the others are worse than 50% error.

6-layers ←

Well, i'm a bit tired of this topic right now.

Let's wrap the findings up, quickly:

• a 6-layer network can solve the 5-operations task, if we add self-attention to the last layer.
• The number of model parameters increased from 237,952 to 389,568 (because i still used the 13-kernels from the 3-layers experiments and the one 64-128 convolution for the self-invented QK)
• The training time increased from 14 to 17 minutes.
• Validation sample error decreased from 99% to 0.06%
• Is this a passable building block for a Very Small Language Model? I might have said yes but right now i say: Don't know, have to test it...

A little graphical disorder at the end of the day: