LoRA with context: a trigram language model

The earlier demo predicted from only the last word. Here the model looks at the last two words: each row of W is a pair (previous word, current word), and each column is a possible next word. So “the ship” and “sea ship” are different rows and can behave differently. LoRA still adds one small correction ΔW = B·A to the frozen W.

Set up

Dictionary: 3 to 5 words, separated by commas. Presets (they reset the words too):

Adapter rank r = 1 (max rank of ΔW) Learning rate = 1.0

Loss (avg KL to target): –

Trainable numbers: –

rank(ΔW): –

rank of the change needed: –

Words whose next word depends on the previous word in the target: –

What the model predicts

Each row is a two-word context. Bars show the probability of each next word, in dictionary order. The orange bar is the word you’re teaching it to say. Scroll inside the table to see all rows.

Original:

Adapted:

Target:

Each line starts from the first two dictionary words and then repeatedly follows the most likely next word, using the last two words as context.

What the adapter is made of

Each column of B times its row of A is one simple “layer” of change; the layers add up to ΔW. You can edit B and A by hand (dashed cells).

In these grids each row is a two-word context (previous, current) and each column is a next word. Orange adds to a score, blue subtracts.

Things to try

With “One context changed”, only the row “sea ship” is taught something new, while “the ship” keeps its old behavior. A bigram model, which sees only “ship”, couldn’t do that, and the page lists “ship” as needing context. The rank you need is about the number of different words you teach, minus one in the extreme case where every row is shifted: softmax ignores adding the same amount to a whole row, so the page measures the needed rank after centering each row. Try “Echo the previous word” with r too low, then raise r.