This blogpost was created as a part of the AI Safety Fundamentals course by BlueDot Impact. All of the code can be found on my GitHub.

TLDR: I tried to uncover if there are specific components in language models that enable typo correction. I identified a subword merging head in the first layer of Llama-3.2-1B that plays a crucial role in the process. In this blog post, I’ll walk through my thought process and findings.

Motivation

Large language models are getting significantly more capable every month, but we still don’t know how they work inside. If a model generates an incorrect answer, we currently have (almost) no way of explaining why it did so. Mechanistic interpretability tries to solve this problem. And while we probably won’t be able to completely reverse engineer large language models anytime soon, interpreting smaller parts of the models is still very valuable since these findings often generalize to larger models and provide valuable insights into how they work.

I was reading a lot of mechanistic interpretability papers over the past year, but I have never actually done anything myself. So I finally gave it a try and explored how typo correction works in LLMs.

Initial exploration

I started by manually testing sentences with typos and looking at their attention patterns. I suspected an early attention head might specialize in typo correction. And sure enough, I found one, though it wasn’t exclusively dedicated to typo correction. I found a subword merging head^[1]. It moves information between tokens that belong to the same word, but were tokenized into multiple tokens. This happens when a word is rare and lacks a dedicated token or, as in my case, when a typo alters its structure.

Attention head ablation

To test whether this head played a role in typo reconstruction, I tried generating text with this head ablated (I zeroed out its output during the forward pass). Without ablation, when given the prompt:

"<|begin_of_text|>What is the meaning of the word 'comptuer'? It means"

the model completed the sentence correctly:

"What is the meaning of the word 'comptuer'? It means a computer. It is a machine that can do calculations."

But when I ablated the subword merging head, the model completely failed:

"What is the meaning of the word 'comptuer'? It means to be a part of something. It is a verb."

Interestingly, ablating this head had no noticeable effect on te...

Posts

Wikitag Contributions

Comments