LLM powered Automatic Text Data Janitor from Microsoft, named Rho-1

A LLM cleaning data of another LLM

Recently a team from microsoft research published a paper titled “Not all tokens are what you need”. Wherein they discuss a approach to automatically detect noisy tokens in a dataset and make the training process ignore them so as to smoothen and speedup training speed.

The below image a visualization of the LLM ignoring attempts to predict such tokens while training.

Illustration of loss removal for undesired tokens from the authors of Rho-1

As you can see in the above image, the highlighted tokens in red dont make any sense. So the LLM also would face a hard time in predicting such tokens. So we simply ask the llm to ignore attempts at learning to predict such tokens by removing the loss component involved for those tokens alone.

But cant we simply remove those tokens from the dataset?

If we remove the tokens from the dataset then the model wont be able to predict sensible tokens when noise is present in the context. The distribution changes by a great ordeal. This is why we simply ask model to prevent attempts at predicting such nonsensical tokens yet have them as noise to predict future sensible tokens.

But how on earth do they detect such tokens?

We first take a LLM which is not finetuned on the target dataset, then compute the losses for each token in the dataset and store them (prefinetuning loss). Then we finetune the LLM on the target dataset and store the new loss values for each tokens as well (post finetuning loss)

Now if the pre-finetuning loss of a token is low and is close to post-finetuning loss, then it means that the original LLM itself had knowledge of predicting such tokens let us call this scenario Low loss (L) -> Low loss(L) i.e L->L.

If the pre-finetuning loss of a token is high and the post-finetuning loss is low, then it means that the tokens follow a pattern which the original LLM didnt learn and that the pattern has been learnt post finetuning. These are the ideal tokens we want to target. Lets call this scenario High Loss -> Low Loss i.e H->L

If the prefinetuning loss of a token is high and the post-finetuning loss is also high, then it means that the tokens likely dont follow a pattern and more random in nature i.e noisy tokens. Lets call this scenario High Loss -> High Loss i.e H->H

There is also a very wierd edge case scenario where the prefinetuning loss of a token in low and the prefinetuning loss is high, meaning the LLM had potentially unlearnt the patterns of such tokens or that the pattern learnt for L->H might have started to over-influence the model. Lets call this scenario Low Loss -> High Loss i.e L->H

The good quality tokens the from the dataset are often the ones in the H->L states. If we focus on such tokens alone then the training loss function will become more smoother thereby enabling faster convergence. And more perfomant model

We filter the H->L tokens by taking the top K% of tokens from the token scores. The token scores is computed by prefinetuning loss — postfinetuning loss

The selected tokens becomes the desired tokens, the remaining becomes undesidered tokens like the ones show in the above example diagram

Isnt it mindblowing?

[embed]

This reminds me of anomaly detection in machine learning where a model learns to take an input encode it and decode it, if the real world input is different from the dataset, the model’s decoding will be different from that of the input. This is a really brilliant implementation of the same but for a pure decoder only LLM.

Yeah, its all fancy, but does it work?

Diagram from the authors of Rho-1 showcasing the results

From the above diagram you can see that the model with the above mentioned technique ( in red) was able to learn faster and more perfomant than vanilla llm training (blue).

It works beautifully

As an AI engineer I have only one thing to say to the authors of the paper

[embed]

If you like my article give me a clap !! Follow me on Twitter / X and LinkedIN