LoRA on attention: a tiny transformer-style model

This model reads the last three words and decides how much attention to pay to each before predicting the next word. It is first trained (“pretrained”) on your original phrase, then frozen. LoRA adapters on its attention matrices are trained to make it write your target phrase instead.

1 · Your phrases

Original phrase Target phrase (same words, new order)

Examples:

Dictionary (from the original phrase):

2 · Train the adapters

Adapter rank r = 2 Learning rate = 0.010

Put LoRA adapters on: Query (Wq) Key (Wk) Value (Wv) Output (Wo)

Loss: –

Steps: 0

Trainable numbers: –

3 · What the model writes

Each model starts from nothing but a “start” marker and keeps adding its most likely next word.

Original model:

Adapted model:

Target phrase:

Where the model looks

For each context (the last three words, ⟨s⟩ = start), the bars show how much attention the model pays to each of those three words, oldest on the left, with the percentage underneath. The prediction is the most likely next word.

4 · The adapters

Each enabled attention matrix W stays frozen. Its adapter adds 2·B·A. You can edit B and A by hand (dashed cells) and watch the attention bars and the text change.

Things to try

Uncheck Value and Output and train with only Query and Key. The adapter can change where the model looks, but not what it reads, and it usually can’t learn the new order. Then switch the others back on. In real LLMs, LoRA is commonly applied to exactly these attention matrices. Training can occasionally get stuck; press “Reset adapters” for a fresh random start.