Reviewing-LLMs

March 29, 2026 • 2043 words

Reviewing LLMs

I haven't written tech posts for around a whole year. I had a messy divorce with tech. But now I got Master's offers, and I am graduating, hopefully things will get better.

We are going to take a second to review what are LLMs, or what we learned in CS182 and CS285. Here is the Karpathy video: Let's build GPT: from scratch, in code, spelled out and this colab. Honestly, the video looks a bit overhyped with 7 million views. Anyway, we are opening the Colab again.

Review

Let's say, the input is

Batch: 1 (ignore for now)
Input Tokens: T
Embedding Dimension: D
Tokenizer Size: V

1. Input

Tokens [T], Crop if too long (Each token is a one-hot vector, so it should be [T, V], but in practice we use [T] and do a lookup on the Embedding Matrix for the TokenEmbeddings)
Embeddings [T, D] $X = \text{TokenEmbeddings} + \text{PositionalEmbeddings}$ (Embedding Matrix [V, D])
Note: Seems to me like a very debatable practice, we don't need a multiplier to PositionalEmbeddings because TokenEmbeddings are learned, but there seem to be many other justifiable ways to pass the positions

2. Transformer

Part A: Attention (N blocks of (Attention Divided in H + Feed Forward)) [T, D]

LayerNorm [T, D] applied across rows, such as $A_{11}, A_{12}... A_{1n}$ (Usually BatchNorm in CNN, here we have different input lengths so we can't use BatchNorm. LayerNorm in Transformers, LayerNorm is not "Data Normalization", it throws away the Magnitude)
Note: That's what I learned in "classes", but there is also RMSNorm here. It also seems debatable about which norms to use
$QKV$ :
- $Q = X_{norm} \cdot W_Q$ [T, D]
- $K = X_{norm} \cdot W_K$ [T, D]
- $V = X_{norm} \cdot W_V$ [T, D]
- $W_Q, W_K, W_V$ [D, D]
- Split into H: $Q_i, K_i, V_i, i= 1, 2, ...H$ [T, D/H].
  _Note: I heard LLMs started using the same K matrices for every $K_{i}$ , same V matrices for every $V_{i}$ ._
Attention [T, T]*H: $\text{head}_{i} = \text{Softmax}\left(\text{Mask}\left(\frac{Q_i \cdot K_i^T}{\sqrt{d_k}}\right)\right) \cdot V_i$
- Attention [T, D/H]*H: $S_{i} = \frac{Q_i \cdot K_i^T}{\sqrt{D/H}}$
  Note: We use ${\sqrt{D/H}}$ for variance now, but it seems to me like a very debatable practice where we can use other ways to reduce variance, such as another LayerNorm
- Cause-al Masking [T, D/H]*H: All values above the diagonal become $-\infty$ , values on the diagnoal or below are unchanged (not "casual" but Cause-al)
- Softmax [T, D/H]*H: Applied across the rows ( $-\infty$ from masking become 0).
- Head Out [T, D/H]*H: $\text{head}\_i = \text{Softmax} \cdot V*{i}$
  _Note: This head practice seems debatable. Maybe we can omit some heads. Maybe we can use more heads than D/H and downscale it, whatever*
AttentionOut:
- $\text {Head} = \text{Concat}(\text{head}\_1, \dots, \text{head}\_H)$ [T, D]
- Weight Matrix $W_O$ [D, D]
- AttentionOut [T, D] $\text{AttentionOut} = \text{Head} \cdot W_O$ [T, D]
Add Residual [T, D]

Part B: Feed Forward (N blocks)

LayerNorm [T, D] applied across rows
Up Projection [T, 4D], Activation (e.g., ReLU) [T, 4D], Down Projection [T, D]
Note: Also seems to me like a highly debatable practice
Add Residual [T, D]

Note: I think the N blocks of same dimension [T, D], where each block is Attention + FeedForward, is debatable. Perhaps we can use several Attentions in one block, drop the FeedForward, use several FeedForward, add something else, or make each block non-identical. Or we can change/shrink the dimension D in the middle.

3. Output

LayerNorm [T, D] applied across rows
Logits [T, V]: $\text{Logits} = X \cdot W_{embed}^T$
Note: We can use a new matrix, but conveniently use dot product with the learned embedding matrix, debatable here
Inference:
- Get last vector $\text{Logits}_{T-1}$ [V] (because every vector can't see the next during training)
- Apply softmax [V]
- Predict next token based on argmax/sampling [1] (here is where all the weird API parameters such as temperature, top-k, top-p come in)
Loss (Training): Shift one (because of the Cause-al Masking). Compare:
- $\text{Logits}_0$ [V] predicts $\text{Token}_1$ [1]
- $\text{Logits}_1$ [V] predicts $\text{Token}_2$ [1]
- ...
- $\text{Logits}_{T-2}$ [V] predicts $\text{Token}_{T-1}$ [1] We use 0-based numbering here
- $\text{Logits}_{T-1}$ [V] predicts the next token, we can grab T+1 during training
- For each of the $T-1$ vectors [V], apply softmax for distributions and calculate the Cross-Entropy loss. Then average it.
  Note: Calculating the loss this way is another seemingly debatable practice

So I learned that almost every "debatable" point here had already turned into several papers in top conferences.

What Exactly Gets Trained?

It is a legit question.

Input: Embedding Matrix [V, D], (maybe) Positional Embeddings [T, D]
Transformer (N blocks):
- LayerNorm1: Scale [D], Shift [D]
- Attention: $QKVO$ [D, D], (maybe) Bias [D] _(sometimes we share the same Key $K_{i}$ or Value $V_{i}$ )_
- LayerNorm2: Scale [D], Shift [D]
- Feed-Forward:
  - Up-projection $W_{1}$ [D, 4*D], (maybe) Bias [4*D]
  - Down-projection $W_{2}$ [4*D, D] (maybe) Bias [D]
Output:
- LayerNorm: Scale [D], Shift [D]
- (maybe, or just use Transpose of Embedding Matrix) Un-embedding Matrix [D, V]

Some More Random Thoughts

What happens if you use pre-trained token-level embeddings, freeze them, and train the Transformer (QKV, etc) from scratch?
Is it a bad idea to use very short input sequences to predict the next token? E.g. what follows "I"? (like, am, do, don't, should, etc) Should we only calculate the loss when the input token is at least a certain number?
Perhaps the token prediction is why sometimes LLMs spit out a word in a different language. Should I force it to spit out only the same alphabet?

Some Cringe

Andrew Nj talked about something very misleading called "Data Normalization" for Housing Prices in his ML Coursera course, which saves the Mean and Standard Deviation, then rescales it back to the output. As it was my first course in LLM 3 years ago, the idea of rescaling back to the output got stuck there, and I ended up confusing this with "LayerNorm" (which doesn't rescale it back) for a really long time.
Then 2 years ago when I was taking CS182 in Cal, I tried to understand the math in the original Attention paper as absolute truths like a math theory. But nobody justified why exactly it they did that. I had to memorize the formulas and the dimensions, which I then forgot quickly. I only got a B+. Now I see these steps as not absolute truths but actually very ambiguous and questionable.
Because I always started to mess up the dimensions once it gets into the Transformer, I did not understand what "Multi-Head Attention" means. I thought "Multi" was trivial and it was just applying the Attention to every Token Embedding. Then I thought "Multi" was applying several parallel, equally large [T, D] $QKV$ and then shrinking it back. In fact, it meant dividing the matrix up on the Embedding Size D, calculating with smaller sizes, and concatenating it back together to the same size.

Kobzar (Taras Shevchenko)

Source Code: https://github.com/jimchen2/kobzar-gpt/

Here is the txt: https://royallib.com/book/shevchenko_taras/kobzar.html

I cleaned the front and the end text.

!wget https://raw.githubusercontent.com/jimchen2/kobzar-gpt/refs/heads/main/kobzar.txt
with open('kobzar.txt', 'r', encoding='utf-8') as f:
    for _ in range(3):
        print(f.readline(), end='')

We take a second to appreciate the poems of Taras Hryhorovych.

Причинна

Реве та стогне Дніпр широкий, Сердитий вітер завива, Додолу верби гне високі, Горами хвилю підійма. І блідий місяць на ту пору Із хмари де-де виглядав, Неначе човен в синім морі, То виринав, то потопав. Ще треті півні не співали, Ніхто нігде не гомонів, Сичі в гаю перекликались, Та ясен раз у раз скрипів.

Tokenizing

Using "o200k_base"

TOKEN ID   | TEXT CHUNK
------------------------------
4647       | 'Р'
4333       | 'ев'
327        | 'е'
6572       | ' та'
3242       | ' ст'
1443       | 'ог'
3907       | 'не'
5117       | ' Д'
6774       | 'ні'
9501       | 'пр'
187104     | ' широкий'

Using "Meta-Llama-3-8B"

TOKEN ID   | TEXT CHUNK
------------------------------
34604      | 'Р'
33742      | 'ев'
1532       | 'е'
100344     | ' та'
18868      | ' ст'
14837      | 'ог'
79862      | 'не'
39903      | ' Д'
123058     | 'ніп'
2233       | 'р'
112781     | ' широк'
44938      | 'ий'

I think for a tokenizer that divides up to alphabets, it is the same for Russian and Ukrainian except for chunks with "і", "ї", "є". ("ґ" is rarely used)

I am just using o200k_base for now.

Vocabulary Size (V): 200019
Total tokens in Kobzar: 209951

Training

We use the classic transformer with some basic parameters so that we can do it on Colab

B = 8          # Batch Size
T = 32         # Context length (Input Tokens)
D = 256        # Embedding Dimension
H = 8          # Number of Attention Heads
N = 16         # Number of Transformer Blocks
MAX_ITERS = 3000
TOKENIZER = o200k_base

We are overtraining a little

Step 0: Loss = 12.4323
Step 500: Loss = 5.3305
Step 1000: Loss = 5.2111
Step 1500: Loss = 4.9537
Step 2000: Loss = 4.5327
Step 2500: Loss = 3.7742
Step 2900: Loss = 3.9520

Generating

The model is 428 MB. It's easy to generate so I did not upload it to HF.

Empty prompt

! Не знову розліпала… за не тілько колись І згадаю сіли що. Як лист, що ледащо царі очі свої діти, І тліс частом Лебирали.

Prompt: 'Як умру'

Як умрусь старим.

Згадаюся; гетьманіти гарматі, Препоровля козачі Петра розвіли добали, Під сонце падалоскіленьку Дону?” Із гор

Prompt: 'Україна'

Україна, У лич самім світі жить!

(Гай при декабря 1844 року, февраля Оні декабря г а

У Києенькаєте, Переолись на виглядаять, Ще ви;ночі

Prompt: 'Кайдан'

Кайданним, У школідуть, Дніпра в княохні стоять, чи вигляда на гори? Зайточко за ними на папи й ночі - Дивитись, пишнійчиком,

Prompt: 'Москва'

Москваь Мар'яні годувала, Аж у хаті зустріну гончосар, А ти не гляне в Київ,- Долітром, Бо за ставами; Зотрі таячи, Не пок

Prompt: 'Місяць'

Місяцькому бні похилилась; Півувались душу хочеться втікая живим На широкій ставом Стебели, І знай на вряду Хвала зелені правдеш!.. гу. Вранцілуери

Prompt: 'Мороз и солнце; день чудесный'

Мороз и солнце; день чудесный, Своєю, братія,- хочеш,- пели синє! Байнуть регезапно, це взяти… бо тим часом своїм тихо він здається. І судить, сохах курейських не дознал

Prompt: 'Шевченко'

Шевченко, з людьми, ставом: «Ніхів лютим». «І тобі часточку,-…» - стара гуляють,- є коси Всь кесаря,- свекій єдиная дівчата… Та,

Prompt: 'Козак і дівчина'

Козак і дівчина його троє.

[Друга половина 1848, Косарал]

йте, Моя чика мати виці колись б стали, а без світі. Прозівну, То б'ї незднімі

Prompt: 'іде'

іде перший! Коло днібити - Зсипаютьаторіміє дніїмурити, немає. Гулянь і Занті! Ярема згадаєшся я возами І гетьманом

Prompt: 'Ой у лузі червона калина,'

Ой у лузі червона калина, І нажнуї білії в за молодій світає, Мовпа варить серденько за хвил ноября катами, Що про ралом, показіє. Та знаєш громаденько дивував

Prompt: 'Цвіте терен, цвіте терен,'

Цвіте терен, цвіте терен, покопилась І за воду Нема сім світі люде. нема кому згадаю. Аж смотреньку дивились, як Алківа, Аж немає, А святая плакала.

Prompt: 'Ой там на горе! Ой там -- на крутой,'

Ой там на горе! Ой там -- на крутой, неді прокн вийми, погані, І не схила! Ще сама на світ батьку благала, і тимкою в не сушіки”. Аж там той Фним Злялине слово.

Prompt: 'Щедрик,'

Щедрик, гарнесенько, й пропила, А в слова тином. Довгофеньки Іншимиікевне півчата пива. І капають, в лісі материїтьтесь І думали С кращ

Prompt: 'Днепр'

Днепріп дитину, Зіброві розмовляла. Я їх не заховали бий, дні ляхам, бачій, От бачрод, У Диві-червоніє В нікому Заховали

Prompt: 'Горя и моря,'

Горя и моря, апалдні брового, А думка, бо не почує тиждень ясний. Гом діялось… взялось взняя�ло. І заспіти одна, не в покривають люди

Chinese Prompt (should be random): '我在那里,'

我在那里, тихо. «Щеженітлось! Тату!” І нищем'юки найми Та й похили. Добустили. Чого то зо благай довгомаха в п'ю Та не найде

(Again)

我在那里, За татарі сестро! Р самыйлу дав мені До улиані, мій божий? Бо арпіка серця, бо й не питрипирокриває, Щоб не буде, Чург на чуж

English Prompt (should be random): 'Once upon a time'

Once upon a time в дорогу хату. Тілько шкурує, Щоб до Гадиંબом у мене трохи, Воскати. Як билиниці робенькі упований, якогось дітатрусу,- Хоч ї

English Prompt (should be random): 'Artificial intelligence will eventually'

Artificial intelligence will eventually, щоб мене не спочило, Не жив, битину В поліму погадне: бодай, що лютно подивилась, Сироти, мов - Марусь: той… отамани брати - До шинку поз

Loading comments...