Experiments

8 min read

LLM Architecture

Transformers

Transformer Types

Encoders

  • Takes text (or other data) as input and outputs a dense representation (or embedding) of that text
  • Example: BERT from Google
  • Use Cases: Text classification, semantic search, Named Entity Recognition
  • Typical Size: Millions of parameters

Decoders

  • Focuses on generating new tokens to complete a sequence, one token at a time
  • Example: Llama from Meta
  • Use Cases: Text generation, chatbots, code generation
  • Typical Size: Billions of parameters

Seq2Seq (Encoder–Decoder)

  • Combines an encoder and a decoder
  • Encoder first processes the input sequence into a context representation, then the decoder generates an output sequence.
  • Example: T5, BART, Whisper
  • Use Cases: Translation, Summarization, Paraphrasing
  • Typical Size: Millions of parameters

Context Length

  • Context length refers to the maximum number of tokens the model can consider at once when generating a response
  • It represents the maximum attention span it has
  • gemma4:26B has context length of 256K tokens

Sliding Window Size

  • Maximum number of historical tokens a model’s local attention layer can look back on at any given moment
  • gemma4:26B has sliding window size of 1024 tokens

Vocabulary Size

  • Total number of unique words or tokens your model knows
  • gemma4:26B has vocabulary size of 262K

Embedding Size

  • Number of dimensions used to mathematically represent each token
  • embeddinggemma has 768 embedding length

Parameters

  • The total count of weights and biases in the neural network, representing the model’s overall memory capacity
  • gemma3:4b has 4.3B parameters

Layers

  • A layer is a component of a neural network that performs a particular computation
  • Model depth generally refers to the number of Hidden Layers
    • Generally Hidden layers refer to Transformer blocks
  • Types of Layers
    • Embedding Layer
      • Common Identifier: embed_tokens
      • Converts token IDs into learned vectors (embeddings)
    • Hidden Layer
      • aka Transformer block
      • Common Identifier: model.layers.N
      • Processes hidden representations using attention and feed-forward networks
    • Normalization Layer
      • Common Identifier: norm
      • Normalizes activations
    • Output Projection layer
      • Common Identifier: lm_head
      • Projects hidden representations to vocabulary logits for next-token prediction
    • Specialized layers
      • For example, vision encoders in multimodal models
  • Examples

Logit

Binary Classification

Multi-Class Classification

  • use softmax function
  • https://en.wikipedia.org/wiki/Softmax_function
  • used when 3 or more mutually exclusive classes
  • In an LLM, each class corresponds to a possible token in the model’s vocabulary
  • All the probabilites obtained by applying softmax function add up to 1
  • be logits for classes

Key-Value Cache

  • aka KV Cache
  • Stores information about previously processed tokens
  • During generation:
    • First pass: Processes all input tokens and stores their information in the cache
    • Next passes: Processes only the newly generated token instead of resending all previous tokens
  • It speeds up text generation by avoiding redundant computation
  • KV cache vs. session:
    • KV cache stores temporary computation information for a sequence
    • A session maintains conversation state, such as chat history
    • The KV cache is typically reused within a generation sequence, but may also be retained or reused across requests by some inference systems

Attention

  • It is process of identifying the most relevant words to predict the next token
  • Intuition: Query and key determine where to pay attention; values determine what information is retrieved
    • — Query
      • Represents what the current token is looking for in other tokens
    • — Key
      • Represents what each token offers for matching against a query
      • Retrieved from the KV cache
    • — Value
      • Represents the information contributed by each token when it receives attention
      • Retrieved from the KV cache
    • — Attention scores
      • Measures how strongly each query matches each key
    • — Key dimension
      • The number of dimensions in each key vector
    • — Scaling factor
      • Prevents attention scores from becoming excessively large.
    • Softmax: Converts scores into attention weights that sum to 1 across the attended positions
    • Multiplication by : Computes a weighted combination of value vectors, producing the attention output

Mixture of Experts (MoE)

LLM Model Contents

  • Example:
  • Architecture and Weights
    • config.json: number of layers, attention heads, hidden layers, base model type
    • .safetensors: weights of model
    • .safetensors.index.json: map of each layer to weights file, used for big models split into multiple weights file
  • Tokenizer
    • tokenizer.json: Vocabulary of model mapping each word/syllable/character to unique ID
    • tokenizer_config.json: Tells library how to handle special tokens and rules like max sequence length etc.
    • special_tokens_map.json: explicitly defines which IDs represent structural markers like the start, end, or padding of text
  • Chat
    • chat_template.jinja: Template to structure text into raw prompt strings
  • Output Control
    • generation_config.json: default behavior when the model generates text like temperature, top_p, top_k
  • Multimodal
    • preprocessor_config.json: Feature extractor containing mathematical arrays for image resizing, cropping, or audio sample normalization.
      • contains processor class
    • processor_config.json: links tokenization with image/video/audio processing
      • contains processor class
  • LoRA
    • adapter_config.json: tells loading library how to attach the plugin to the base model
    • adapter_model.safetensors: adapter weights

Running Raw LLM Model

  • Example Model: https://huggingface.co/HuggingFaceTB/SmolLM2-360M
  • https://github.com/huggingface/transformers/blob/536ecc007387a50e77603bb5d92100e9b07514cc/src/transformers/models/llama/modeling_llama.py#L356
  • https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_outputs.py
  • https://github.com/huggingface/transformers/blob/09281b8d56b08199ce9f397c694f38853cc45278/src/transformers/modeling_utils.py#L3823
  • For inference model.generate() is called which internally calls forward() method
  • generate() repeatedly calls the model’s forward() method to obtain next-token predictions
  • forward() method performs one pass over all the layers in LLM Neural Network
    • Input
      • input_ids: Tensor
        • Integer IDs representing the input tokens
      • inputs_embeds: Tensor
        • Vector representations of input tokens
        • Pass either input_ids or input_embeds
      • attention_mask: Tensor
        • consist of 0s and 1s
        • Indicates which input positions are valid (1) and which should be masked (0), commonly to exclude padding tokens
      • position_ids: Tensor
        • Specifies the position index of each token in the sequence
        • If KV caching is disabled, positions typically start at 0 for the input sequence.
        • If KV caching is enabled, positions for newly processed tokens typically continue from the number of previously cached tokens.
      • use_cache: bool
        • Controls whether the model uses the key-value cache
      • past_key_values: Cache
        • KV cache
        • Stores previously computed attention keys and values
          so they can be reused instead of recomputed during generation
    • Output
      • logits
        • Logit score for each token in vocabulary
      • past_key_values
        • The updated key-value cache containing attention information 
          for tokens processed so far
        • This is passed into the next forward call to speed up generation
  • In essence, LLM does following:
    • Take input tokens from its vocabulary
    • Converts them into vector representations using embedding layer
    • Processes them through multiple neural network layers using attention and feed-forward operations
    • Each pass produces logits representing the scores of possible next tokens
    • During text generation, these logits are converted into probabilities
    • Using those probabilities and LLM config, a next token is selected
    • This next token is fed back into the model to predict the next token
      • With KV caching enabled, the model typically processes only the new token on subsequent steps
    • Process repeats until a stopping condition is met
from transformers import AutoTokenizer, AutoModelForCausalLM
 
model_name = "HuggingFaceTB/SmolLM2-360M-Instruct"
 
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
 
# 49152 ~ 49k tokens
print("Tokenizer vocabulary:", len(tokenizer)) 
 
# model_type: llama
# dtype: bfloat16
# num_hidden_layers: 32
# hidden_size: 960 (embedding size of each vector)
# num_attention_heads: 15
print("Model config:", model.config.to_json_string())
 
# Prepare a conversation
messages = [
    {"role": "user", "content": "Explain branches of mathematics."}
]
 
# Convert conversation to model input
inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_tensors="pt",
)
 
# torch.Size([1, 35])  
# 1 conversation with 35 tokens  
print(inputs["input_ids"].shape)  
  
# torch.Size([1, 35])  
# consist of 0s and 1s for each input token  
print(inputs["attention_mask"].shape)
 
# Generate text
# generate internally calls `forward()`
outputs = model.generate(
    **inputs,
    max_new_tokens=200,
)
 
# include input tokens (35) + newly generated tokens (200)
print(outputs.shape)
 
# Decode only the newly generated tokens
new_tokens = outputs[0][inputs["input_ids"].shape[-1]:]
response = tokenizer.decode(new_tokens, skip_special_tokens=True)
 
print(response)
  • Model Config
{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 1,
  "dtype": "bfloat16",
  "eos_token_id": 2,
  "head_dim": 64,
  "hidden_act": "silu",
  "hidden_size": 960,
  "initializer_range": 0.02,
  "intermediate_size": 2560,
  "is_llama_config": true,
  "max_position_embeddings": 8192,
  "mlp_bias": false,
  "model_type": "llama",
  "num_attention_heads": 15,
  "num_hidden_layers": 32,
  "num_key_value_heads": 5,
  "pad_token_id": 2,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_interleaved": false,
  "rope_parameters": {
    "rope_theta": 100000,
    "rope_type": "default"
  },
  "tie_word_embeddings": true,
  "transformers.js_config": {
    "kv_cache_dtype": {
      "fp16": "float16",
      "q4f16": "float16"
    }
  },
  "transformers_version": "5.8.1",
  "use_cache": true,
  "vocab_size": 49152
}