Experiment log book

Experiments summary

def collatz(n):
    if n == 1:
        return 1
    elif n % 2 == 0:
        return 1 + collatz(n // 2)
    else:
        return 1 + collatz(3*n + 1)

Applying dropout in the attention on the tokens before the closest [compute]. For instance, if we have ["enxhell" "is" "[compute]" "[compute]" "18"]. When outputing "18" the attention module can only attend to the closest compute and the rest of the tokens are randomly dropped out heavily (0.6+ probability of dropout). The second compute token can only attend at the previous compute token (that's the closest) and the rest are heavily dropped out. This way forcing the model to store as much as possible in the compute tokens. This did not improve or reduce the performance. Which is perhaps a sign that the model is already utilizing the compute tokens - it is storing something useful there, otherwise the model is almost blind to the input. Even if it's a summary of the previous tokens instead of deciding to take the easy way out and predicting "[compute]" without doing anything useful. Although, the model could in theory do another degenerate case of simply predicting the output in the first [compute] and then copying it over and over.
Make sure that in the output of the first compute token we try to predict the first non-compute token that follows the compute block. So in the case above, when we calculate the loss on the first "[compute]" token output we combine the loss of predicting the next token (in this case "[compute]" and the first non-compute token which is "18"). So that way the [compute] token is forced to do something useful. Additionally, the loss is weighted by how far the token is. Specifically, 1/distance_in_tokens. So for instance for the first compute, the loss is 1/2 and for the second compute it's 1/1 = 1. This also did not impact performance.



Performance of the models of various depths and IC tokens grouped by the number of operations required to output that digit correctly

Log Book



def collatz(n):
    if n == 1:
        return 1
    elif n % 2 == 0:
        return 1 + collatz(n // 2)
    else:
        return 1 + collatz(3*n + 1)



Varying number of compute tokens and depth

{n}---[compute][compute]...[compute]{collatz(n)}

The depth doesn't matter
Number of compute tokens doesn't matter either. In fact, when we add 50 compute tokens, the model destabilizes a bit. But in essence the performance is roughly the same in evaluation accuracy. MSE loss also shows a similar story.
Train/loss however does show that a deeper model is able to fit the data more. Note that for this, a fair comparison is only between the same compute tokens, but different depths. Because if you compare say 1-ic and 5-ic the loss for 5-ic is going to be lower because it's measured across all tokens and the IC tokens are easy to predict. However, the lower loss doesn't seem to materialize to better eval

The model is not able to utilize the compute/IC tokens.
Given that the depth of the model didn't help, maybe the task is just too hard.

Some caveats:

Tokenization is standard gp2 BPE. So trying to predict the original tokens has one more possibility of failure - i.e some tokens like 42 is one token, numbers like 139 are two, but some three digits are 1 token and so on and so forth. This could maybe make the task a bit too hard for a small model and just adds one more variable that's not needed. I will run experiments without this though, more on this later.
Evaluation is done using teacher-forcing by hugging face. That is a fancy way of saying that for each token the prediction is done such that the model has the ground truth for all the previous tokens before that. So, it is not generating the whole sequence, but only the next token. So if say the output is 139 and the tokens are "1" and "39" then when it's predicting "39" it already knows that the first token is "1", even if it messed up before. Now that's not a big deal for accuracy because if it messes up one token, it's wrong already. It should only affect the MSE and if we down the line measure accuracy per digit. Will correct this going forward

Forcing the model to use [compute]

Redoing the experiments but with single digits as tokens and without teacher-forcing

Make the tokenizer tokenize each digit separately. So the vocab of GPT to be single digits, and some basic signs
Make the evaluation use generation instead of teacher-forcing.

Does the task benefit from more compute?

3-digit multiplication







Accuracy at each digit level. Digit positions are measured from right to left, so 6 is the hundreds of thousands digit.









from collections import defaultdict
def count_operations_per_digit_with_carries(a, b):
    # a and b are two 3-digit numbers
    additions = defaultdict(int)
    multiplications = defaultdict(int)
    carry_additions = defaultdict(int)

    for i, digit_a in enumerate(reversed(str(a))):
        for j, digit_b in enumerate(reversed(str(b))):
            position = i + j + 1
            multiplications[position] += 1
            if multiplications[position] > 1:
                additions[position] += 1
            product = int(digit_a) * int(digit_b)
            carry = product // 10
            if carry > 0:
                carry_additions[position + 1] += 1

    for position, count in carry_additions.items():
        additions[position] += count

    return additions, multiplications





Do more tokens help? Tending towards yes, but the relationship is not as smooth as more-tokens better-performance. Also only models of certain size seem to be able to utilize them (depth 1 can't) and beyond depth-6 in our case where the model achieved the ceiling, it was counterproductive.
Why is there variance across runs (answer below)? I don't really have an answer why there is variance between runs. Perhaps the model has trouble always effectively using the extra computation / registers? Or maybe it's a problem with dataset change. We'll see once those experiments finish running.
What are these models learning? They seem to be progressively closing the gap with more and more operations as seen by the performance in the operations. Digit 3 and 4 are the most problematic because they require the most operations. Which would align with the idea that compute tokens could be useful to store intermediate computations.

The variance between different runs

Dataset was slightly different as I was sampling random pairs but the seed wasn't set. Although I was sampling a lot of data so that wouldn't make a big difference. That would also result in a different test set though and that maybe might make a difference
Initialization of the network

Update - Related Work