The AiEdge Newsletter

AdalFlow: A PyTorch-Like Framework to Auto-Optimizing Prompt for your LLM agent

Damien Benveniste — Mon, 29 Sep 2025 15:01:00 GMT

AI Agent frameworks are becoming just as important as model training itself! I am excited to introduce you to Li Yin. She is the CEO of SylphAI and the founder of AdalFlow, a PyTorch-like open-source library on GitHub that enables developers to build and auto-optimize any Language Model (LM) workflows.

In this guest post, Aria Shi, the Developer Relations lead at SylphAI, walks you through how AdalFlow empowers AI Agent development, highlighting a hands-on example with a LinkedIn Reachout Agent.

Say goodbye to manual prompt engineering. AdalFlow is the all-in-one, auto-differentiative solution for optimizing prompts, whether you’re using zero-shot or few-shot learning. Backed by our state-of-the-art research (LLM-AutoDiff and Learn-to-Reason), our framework achieves the highest accuracy among all automatic prompt optimization libraries.

The rise of large language models has completely changed the way we build applications—whether it’s chatbots, RAG systems, or fully autonomous agents. But as an AI engineer, trying to bring these models into production often feels like stitching together a bunch of experiments, rather than building a stable and reliable system.

We introduce AdalFlow: a PyTorch-like library designed to bring structure, clarity, and optimization to the world of LLM application development. Built as a community-driven project, AdalFlow is uniting AI research and production engineering into a single ecosystem.

AdalFlow GitHub Repository

Why We Built AdalFlow

Modern AI development faces a paradox. On one hand, researchers push the boundaries of model capabilities with new techniques in prompting, evaluation, and optimization. On the other hand, production teams need reproducibility, scalability, and a way to iterate safely on real-world data.

Most libraries excel at one side of the equation but leave the other underserved. AdalFlow was born to bridge this gap. With 100% control and clarity of source code, it empowers researchers to experiment freely while giving product engineers the tools to build and ship with confidence.

Why AdalFlow Matters

By treating prompts as first-class citizens and introducing LLM-AutoDiff, AdalFlow provides what’s been missing in the LLM ecosystem:

For researchers: A familiar PyTorch-like environment to prototype new prompting and training methods.
For engineers: Production-ready workflows that are debuggable, reproducible, and optimizable.
For teams: A shared framework that unites research and production into one healthy ecosystem.

The overview of AdalFlow

Core Philosophy: Prompt Is the New Programming Language

If PyTorch turned tensors into the lingua franca of deep learning, AdalFlow treats prompts as the new programming primitives.

Every LLM application boils down to structured prompts and their transformations. AdalFlow embraces this reality by making prompt engineering explicit and optimizable. Behind the scenes, it uses the Jinja2 templating engine to let developers define composable prompt structures, ensuring that LLM apps are both modular and debuggable.

Components: The Building Blocks of LLM Workflows

At the heart of AdalFlow lies the Component abstraction. Just as nn.Module became the foundation for PyTorch models, Components unify every stage of an LLM pipeline.

AdalFlow Core Concepts

Component: The base class for all workflows. Handles both training (forward) and inference (call) modes, with bicall bridging the two.
GradComponent: Components capable of backpropagation (e.g., Generators, Retrievers).
DataComponent: Lightweight components for formatting and parsing data (e.g., DataClassParser).
LossComponent: Wraps evaluation metrics and enables gradient-like feedback for text optimization.

Example 1: Q&A with Object Counting (Component + DataComponent)

Question Answering - Build and Optimize LM Workflows

template = r"""
{{system_prompt}}


{{input_str}}
"""

@adal.func_to_data_component
def parse_integer_answer(answer: str):
    numbers = re.findall(r"\d+", answer)
    return int(numbers[-1])

What’s happening here?

parse_integer_answer is wrapped with @adal.func_to_data_component.
This turns a plain Python function into a DataComponent, which handles structured output parsing.
In this case, it ensures the model’s answer ends with a numerical value.

Next, we define a full pipeline:

class ObjectCountTaskPipeline(adal.Component):
    def __init__(
            self, model_client: adal.ModelClient, model_kwargs: Dict
        ):
        super().__init__()
        system_prompt = adal.Parameter(
            data="You will answer a reasoning question. Think step by step. The last line should be 'Answer: $VALUE'.",
            role_desc="Task instruction for the model",
            requires_opt=True,
            param_type=ParameterType.PROMPT,
        )
        self.llm_counter = adal.Generator(
            model_client=model_client,
            model_kwargs=model_kwargs,
            template=template,
            prompt_kwargs={"system_prompt": system_prompt},
            output_processors=parse_integer_answer,
        )

    def bicall(self, question: str, id: str = None):
        return self.llm_counter(
              prompt_kwargs={"input_str": question}, id=id
        )

ObjectCountTaskPipeline subclasses Component. Inside it, we define:

A Parameter of type PROMPT, which AdalFlow can later auto-optimize.
A Generator (a GradComponent) that executes the prompt, then passes the raw LLM output through our parse_integer_answer DataComponent.

The workflow is:
Prompt → LLM Generation → Structured Output Parsing → Final Numerical Answer.

Example 2: Classification with Structured Output (Component + DataClass)

Classification Optimization - Build and Optimize LM Workflows

Classification tasks are a perfect showcase of AdalFlow’s DataClass feature.

@dataclass
class TRECExtendedData(adal.DataClass):
    question: str = field(
        metadata={"desc": "The question to be classified"}
    )
    rationale: str = field(
        metadata={"desc": "Step-by-step reasoning"}, default=None
    )
    class_name: Literal[
        "ABBR", "ENTY", "DESC", "HUM", "LOC", "NUM"
    ] = field(
        metadata={"desc": "The class name"}, default=None
    )

    __input_fields__ = ["question"]
    __output_fields__ = ["rationale", "class_name"]

TRECExtendedData extends DataClass, which (like Pydantic) gives us schema enforcement.
Input: a question.
Output: a rationale (reasoning trace) and a class_name (final label).

Now let’s plug it into a pipeline:

class TRECClassifierStructuredOutput(adal.Component):
    def __init__(
        self, model_client: adal.ModelClient, model_kwargs: Dict
    ):
        super().__init__()
        # Task description prompt
        task_desc_str = adal.Prompt(
            template=task_desc_template,
            prompt_kwargs={
                "classes": [
                    {"label": l, "desc": d} 
                    for l, d 
                    in zip(_COARSE_LABELS, _COARSE_LABELS_DESC)
                ]
            }
        )()

        parser = adal.DataClassParser(
            data_class=TRECExtendedData,
            return_data_class=True,
            format_type="yaml"
        )

        prompt_kwargs = {
            "system_prompt": adal.Parameter(
                data=task_desc_str,
                role_desc="Task description",
                requires_opt=True,
                param_type=adal.ParameterType.PROMPT,
            ),
            "output_format_str": parser.get_output_format_str(),
        }

        self.llm = adal.Generator(
            model_client=model_client,
            model_kwargs=model_kwargs,
            prompt_kwargs=prompt_kwargs,
            template=template,
            output_processors=parser,
        )

    def bicall(self, question: str, id: Optional[str] = None):
        return self.llm(prompt_kwargs={"input_str": question}, id=id)

The Prompt defines the system instruction with class definitions.
DataClassParser enforces structured YAML output that matches TRECExtendedData.
Generator (a GradComponent) runs the LLM with prompt + parser.
Output is guaranteed to follow the schema: rationale + class name.

This ensures the model never drifts into free-form answers—it always returns structured classification results.

Example 3: Training With LossComponent

Finally, how do we train or optimize these components? That’s where LossComponent comes in:

eval_fn = AnswerMatchAcc(type="exact_match").compute_single_item
loss_fn = adal.EvalFnToTextLoss(
    eval_fn=eval_fn,
    eval_fn_desc="exact_match: 1 if str(y) == str(y_gt) else 0"
)

AnswerMatchAcc is the evaluation metric.
EvalFnToTextLoss wraps it as a LossComponent, enabling LLM-AutoDiff to optimize prompts automatically during training.

By attaching this to your pipeline, you get a full training loop:
Forward pass → Eval metric → Backward engine → Prompt optimization.

Agents: Reasoning Meets Action

AdalFlow embraces the ReAct paradigm—combining reasoning (plan) with acting (tool use)—to build autonomous, auditable AI systems. An agent reasons about the task, selects tools, executes them, observes results, and iterates until it can deliver a final answer.

Architecture at a Glance

Agent (planner + tool manager)
Handles planning and decision-making via a Generator-based planner, and knows what tools are available and how to call them.
Runner (executor + conversation loop)
Orchestrates multi-step execution, tool calling, observation handling, timeouts, and final answer synthesis.

This separation lets you swap or customize planning vs. execution independently.

Execution Flow (ReAct Loop Recap)

Planning – The Agent (Generator planner) analyzes input and proposes the next action.
Tool Selection – Chooses a tool from the registered set.
Tool Execution – The Runner invokes the tool with arguments.
Observation – The result is fed back to the planner.
Iteration – Repeat 1–4 up to max_steps or until confident.
Final Answer – The planner synthesizes the answer (optionally into a structured type).

Minimal, End-to-End Example

1) Define a Tool (callable or FunctionTool)

# Tool: a plain Python callable works, or wrap with FunctionTool for extras.
def calculator(expression: str) -> str:
    """Evaluate a mathematical expression."""
    try:
        result = eval(expression)
        return f"Result: {result}"
    except Exception as e:
        return f"Error: {e}"

2) Build the Agent (Planner + Tools)

from adalflow import Agent, Runner
from adalflow.components.model_client.openai_client import OpenAIClient

agent = Agent(
    name="CalculatorAgent",  # Agent identifier
    tools=[calculator],  # List of tools (callables or FunctionTool)
    # LLM client used by the planner (Generator-based)
    model_client=OpenAIClient(),
    model_kwargs={"model": "gpt-4o", "temperature": 0.3},
    max_steps=6,  # Upper bound for ReAct loops
)

What this maps to:

Planner: An internal Generator that decides the next step (think: “reasoning trace”).
ToolManager: The agent’s registry of permitted tools.
max_steps: Safety rail to prevent runaway loops.

Model Configuration (Swap Backends Easily)

# OpenAI
from adalflow.components.model_client.openai_client import OpenAIClient
agent = Agent(
    model_client=OpenAIClient(), 
    model_kwargs={"model": "gpt-4o"}
)

# Anthropic
from adalflow.components.model_client.anthropic_client import (
    AnthropicAPIClient
)
agent = Agent(
    model_client=AnthropicAPIClient(), 
    model_kwargs={"model": "claude-3-sonnet-20240229"}
)

3) Execute with the Runner (Multi-step Orchestration)

# Manages turns, tool calls, observations, and finalization
runner = Runner(agent=agent) 

result = runner.call(
    prompt_kwargs={"input_str": "Invoke the calculator tool and calculate 15 * 7 + 23"}
)

print(result.answer)
# -> "The result of 15 * 7 + 23 is 128."

RunnerResult schema (returned by Runner.call)

# result has:
# - result.step_history: [StepOutput(...)]  # Each step’s action + observation
# - result.answer: str | structured type     # Final synthesized answer
# - result.error: None | Exception info      # Error if something failed
# - result.ctx: dict | None                  # Optional execution metadata

This is the full ReAct loop in action:
Plan → Select Tool → Execute → Observe → Iterate → Answer.

Advanced Features (Production-Ready)

1) Streaming Execution (Real-Time Updates)

# Pseudocode: actual API may differ slightly in your version.
stream = runner.stream(
    prompt_kwargs={"input_str": "Compute 42 * 73 and explain."}
)
for update in stream:
    # update contains partial thoughts, tool calls, observations, etc.
    print(update)

Use streaming to surface live reasoning/tool progress in UIs.

2) Human-in-the-Loop (Permission Management)

from adalflow.permissions import PermissionManager

class MyPerms(PermissionManager):
    def approve(self, tool_name: str, args: dict) -> bool:
        # Example policy: only allow calculator; prompt user otherwise
        return tool_name == "calculator"

agent = Agent(
    name="GuardedAgent",
    tools=[calculator, search_tool],
    model_client=OpenAIClient(),
    model_kwargs={"model": "gpt-4o"},
    permission_manager=MyPerms(),  # <- Every tool call can be inspected/approved)

Great for tools that hit external systems (files, emails, APIs).

3) Custom System Templates (Planner Behavior)

custom_role_desc = """
You are a careful, step-by-step data analyst.
When you use a tool, explain why and what you expect to get.
"""
agent = Agent(
    name="DataAnalyst",
    # Custom planner persona and guardrails
    role_desc=custom_role_desc,        
    model_client=OpenAIClient(),
    model_kwargs={"model": "gpt-4o", "temperature": 0.2},
)

4) Tracing (Observability)

# Configure tracing once (destination: console, file, or tracing backend)
from adalflow.tracing import enable_tracing
enable_tracing(project="adalflow-agents-demo")

result = runner.call(prompt_kwargs={
    "input_str": "Use the calculator for 88*19."
})
# Inspect step_history, tool IO, latency, errors, etc.

Agent Summary:

Agent = Reasoning + Tool selection (Generator-based planner + ToolManager)
Runner = Controlled execution loop (steps, tools, observations, final answer)
Tools = Safe, permissioned extensions to the agent’s capabilities
Production = Streaming, human approvals, tracing, structured outputs

Real-World Use Case: LinkedIn Recruitment Agent with AdalFlow

Hiring top talent is one of the most resource-intensive parts of building a company. Recruiters spend hours scrolling LinkedIn, opening profiles, copying notes, and crafting outreach messages.

What if we could automate that entire workflow—turning hours of manual searching into minutes of AI-assisted sourcing?

That’s exactly what we built using AdalFlow’s Agent + Runner architecture combined with browser automation via Chrome DevTools Protocol (CDP).

✨ Before vs. After

Traditional Recruiting Workflow (Manual)

❌ BEFORE: 2–3 hours per role
1. Navigate to LinkedIn people search
2. Type in “Product Manager, San Francisco”
3. Scroll endlessly, click into profiles
4. Skim experience, education, skills
5. Take notes in spreadsheets
6. Write & send DMs manually
Automated Workflow with AdalFlow (Agentic)
✅ AFTER: 10 minutes per role
1. Run: linkedin-agent --query “Product Manager” --limit 10
2. Agent plans and executes:
- Smart search strategy
- Extract profiles via browser automation
- Evaluate candidates with scoring models
- Draft personalized outreach messages
3. Get structured output: JSON/CSV with names, titles, LinkedIn URLs, evaluation scores, outreach drafts. Recruiters get to focus on talking to people, not copy-pasting data.

How It Works — Global State Architecture

We structured the solution around a global state shared between tools. Each tool contributes partial data (search results, profiles, evaluations, outreach drafts), which the Agent combines into a full pipeline. Agent combines into a full pipeline.

Full Pipeline of LinkedInAgent

Implementation with AdalFlow

We implemented the LinkedInAgent by encapsulating:

Agent → Planner + Tools (search, extract, evaluate, outreach)
Runner → Execution loop with error handling and logging

class LinkedInAgent:
    """
    LinkedIn recruitment agent powered by AdalFlow.
    - Encapsulates Agent + Runner
    - Provides default recruitment tools
    - Supports both sync call() and async acall()
    """

    def __init__(
        self,
        model_client: Optional[OpenAIClient] = None,
        model_kwargs: Optional[Dict[str, Any]] = None,
        max_steps: Optional[int] = None,
        role_desc: Optional[str] = None,
        **kwargs,
    ):
        # Defaults
        model_client = model_client or OpenAIClient()
        model_kwargs = model_kwargs or {
           "model": "gpt-4o", “temperature”: 0.3
        }
        max_steps = max_steps or 6

        # Recruitment workflow tools
        self.tools = [
            # 1. Search LinkedIn via CDP
            SmartCandidateSearchTool,
            # 2. Extract structured profile data     
            ExtractCandidateProfilesTool, 
            # 3. Score candidates   
            CandidateEvaluationTool,
            # 4. Draft personalized outreach         
            CandidateOutreachGenerationTool,
            # 5. Persist results 
            SaveOutreachResultsTool,         
        ]

        # Agent role description (personality / instructions)
        role_desc = role_desc or "You are a recruitment assistant that sources and evaluates LinkedIn candidates."

        # Initialize Agent + Runner
        self.agent = Agent(
            name="LinkedInRecruiter",
            tools=self.tools,
            model_client=model_client,
            model_kwargs=model_kwargs,
            max_steps=max_steps,
            role_desc=role_desc,
            **kwargs,
        )
        self.runner = Runner(agent=self.agent, max_steps=max_steps)

    def call(
        self, query: str, context: Optional[Dict[str, Any]] = None
    ):
        return self.runner.call(prompt_kwargs={"input_str": query})

    async def acall(
        self, query: str, context: Optional[Dict[str, Any]] = None
    ):
        return await self.runner.acall(
            prompt_kwargs={"input_str": query}
        )

Full Workflow Execution

Here’s how we stitch the agent into a production workflow:

def execute_search_workflow(
    self, progress_tracker=None) -> List[Dict[str, Any]]:
    logger = get_logger()
    logger.set_workflow_context("workflow_main", "initialization")

    log_phase_start("WORKFLOW_START", f"Target: {self.limit} candidates for {self.query} in {self.location}")

    candidates = []
    try:
        log_info("🤖 Initializing LinkedIn agent...")
        agent, user_query = self.initialize_agent()

        if progress_tracker:
            progress_tracker.start_workflow()

        # Run full pipeline
        result = agent.call(query=user_query)
        self._print_agent_execution_steps(result)

        # Collect data from global state
        from ..core.workflow_state import get_complete_workflow_data
        workflow_data = get_complete_workflow_data()

        candidates = self._build_complete_candidate_data(workflow_data)

        log_info(f"✅ Found {len(candidates)} candidates")
        return candidates

    except Exception as e:
        log_error(f"❌ Workflow failed: {e}")
        return candidates

Example Output

After running:

linkedin-agent --query “Product Manager San Francisco” --limit 10

We get structured results like:

[
  {
    “name”: “Alex Chen”,
    “title”: “Senior Product Manager @ Stripe”,
    “location”: “San Francisco Bay Area”,
    “profile_url”: “https://linkedin.com/in/alexchen”,
    “score”: 0.92,
    “outreach_message”: “Hi Alex, I came across your experience at Stripe...”
  },
  {
    “name”: “Maria Lopez”,
    “title”: “PM, Growth @ Airbnb”,
    “location”: “San Francisco Bay Area”,
    “profile_url”: “https://linkedin.com/in/marialopez”,
    “score”: 0.88,
    “outreach_message”: “Hi Maria, your background in growth product design really stood out...”
  }
]

Screenshot of the LinkedInAgent

============================================================
WORKFLOW COMPLETION SUMMARY
============================================================
Success: ✅ Yes
Total Candidates: 2
Duration: 68.8 seconds
Session: 20250909_220617
Log Files:
• Main: logs/workflow_20250909_220617.log
• Debug: logs/debug_20250909_220617.log
• Agent Steps: logs/agent_steps_20250909_220617.log
• Errors: logs/errors_20250909_220617.log
============================================================
🏁 MAIN ✅ COMPLETED - Processed 2 candidates
[RESULTS] ✅ Recruitment workflow completed!
[RESULTS] 📊 Final result: Successfully processed 2 candidates

Looking Ahead

As the LLM landscape evolves, frameworks like AdalFlow will become the backbone of application development. Just as PyTorch accelerated deep learning, AdalFlow has the potential to democratize LLM app building—from chatbots to agents to beyond.

If you’re excited about shaping the future of AI workflows, the project is open-source and community-driven. Whether you’re an AI researcher, product engineer, or just curious about building smarter applications, now’s the time to get involved.

🚀 AdalFlow isn’t just another library. It’s a paradigm shift in how we think about programming with language models.

Last Week to Register: Build Production-Ready Agentic-RAG Applications From Scratch Course!

Damien Benveniste — Tue, 23 Sep 2025 15:02:58 GMT

This is the last week to register for the Build Production-Ready Agentic-RAG Applications From Scratch course! This is a fully hands-on course where we are going to implement step-by-step from scratch a production-ready Agentic-RAG application with LangGraph, FastAPI, and React!

Signup!

What we are going to build

We are going to build a fun web application where we can demonstrate how to orchestrate a robust RAG application using LangGraph, FastAPI, and React. Here is what we are going to build:

A user can pass a GitHub repository URL
The files of the related repository are scraped and indexed in a vector database
Now the code is available for the user to ask questions about.

On the frontend, we will need two main functionalities:

A page where we can input the repository URL and start the crawling and indexing processes:

And a chatbot interface to ask questions about the code in the repository:

On the backend, we will need the related endpoints:

The indexing endpoint will respond to the provided GitHub repository URL and the “crawl“ action to start the crawling and indexing processes.
The chat endpoint that will respond to messages sent by the user from the chatbot interface.

We are going to use the following tools:

React for the frontend
FastAPI for the backend
LangGraph for the agentic orchestration
Pinecone for the vector database
Langsmith for observability
Deploy everything on Google Cloud!

Project-based course

We will focus on building the project from the ground up, as we would on the job. Here is how we are going to structure the project development:

Introduction
- What we want to build
- Setting up the environment
The RAG Application
- The Data Parsing Pipeline
- The Indexing Pipeline
- The Basic RAG Pipeline
- Adding Observability to the Pipeline with Langsmith
- Going Agentic
The Backend Application
- The Indexing API Endpoint
- Adding Memory
- Administering the Database Data
The Frontend Application
- The Indexing Page
- The Chatbot Page
Deploying to GCP

Each session will be a live, hands-on coding session where we are going to implement every component from scratch

Going Agentic

“Agentic” means that we are going to use an LLM as a decision engine to enhance the quality of our pipeline. We will focus on improving the accuracy of the pipeline at the cost of latency and cost, and discuss the opportunities to reduce those induced negative points with small language models and fine-tuning. In the RAG pipeline, we are going to build a subagent for each of the main components:

Intent router: the entry point of the pipeline that will decide if a RAG pipeline is required.
The retriever: The sub-agent that will extract the right data
The generator: The sub-agent that will generate the response to the user

Scaling up

With this course, I want to focus on what we would need to do to deploy the application to 1M users. We will make sure to design every endpoint to be asynchronous, queue the indexing requests, and deploy the application with elastic load balancing to scale the application horizontally.

This is going to be a fun ride! Make sure to join us!

Signup!

The Real-World AI Engineering Roadblocks You Face Today

👋 Prototype → Production Gap — Moving from a notebook demo to a secure, observable, multi-tenant service requires orchestration, evals, guardrails, and ops most teams lack.

👋 “Easy RAG” vs “Reliable RAG” — Anyone can retrieve-then-generate; making answers faithful, fresh, fast, and cost-controlled under real traffic is the hard part.

👋 Framework Overload — The ecosystem is noisy; you need clear criteria (maturity, extensibility, latency, cost) and reference patterns to choose confidently.

👋 It’s Software Engineering First — Success hinges on clean interfaces, tests, typed configs, tracing, CI/CD, and change management—not just prompts and models.

👋 From Laptop to 1M Users — Scaling demands streaming, batching, caching, autoscaling, and SLOs, or your p95 explodes and costs spiral.

How this course will help you

✅ Ship a real Agentic RAG app, not a demo — Stand up an end-to-end stack—LangGraph → FastAPI → React, that runs locally today and deploys via a clean, fork-and-ship monorepo.

✅ Make retrieval dependable, not lucky — Adopt schema-aware chunking, strong dense embeddings with sensible metadata filters, and context packing with citations so answers stay faithful, fresh, and concise.

✅ Harden agentic workflows — Design a typed LangGraph state and build nodes for rewrite → retrieve → rerank → synthesize → cite → safety-check, with retries and timeouts so plans don’t loop or stall.

✅ Scale the experience, not the headaches — Enable server-streaming in FastAPI, cap top-k, trim context budgets, and add early-exit rules; deploy with autoscaling so you can serve real traffic without infra fuss.

✅ See enough to fix things fast — Bake in structured logs (no vendor tracing), per-step timing counters, and UI breadcrumbs/citations to follow query → context → answer and spot common failure patterns quickly.

✅ Choose frameworks with confidence — Follow an opinionated reference architecture plus a simple choice rubric (maturity, extensibility, latency, cost, swap effort) so you know when to stick—and how to swap components without rewrites.

✅ Write maintainable RAG code — Use clean module boundaries (ingest / retrieve / rerank / synthesize), typed configs (Pydantic Settings), and sensible secrets/env management so your team can extend it safely.

You’ll walk away with

✨ A running Agentic RAG app (LangGraph + FastAPI + React) in a fork-and-ship monorepo.

✨ An ingestion/indexing pipeline with metadata, hybrid retrieval, and optional re-ranking.

✨ A chat UI with citations, source previews, and conversation memory that behaves.

✨ Deploy scripts and env templates to go live right after class.

✨ A framework choice memo + adapters to swap models/vector stores without starting over.

Bottom line: this isn’t a vitamin, it’s a blueprint you can put in production.

What you’ll get out of this course

Orchestrate complex RAG pipelines with LangGraph and OpenAI API: Build a typed LangGraph that routes rewrite → retrieve → rerank → synthesize → cite → self-check with retries, timeouts, early-exit rules, and real tool calls, exposed as a clean HTTP API.
Build scalable asynchronous applications with FastAPI: Ship async FastAPI endpoints, well-typed request/response models, input validation, and sensible timeouts, ready to run locally and deploy to production.
Implement chatbot interfaces with React: Create a chat UI that shows citations and source previews, lets users scope queries, preserves safe chat history, and handles transient API errors gracefully.
Mitigate hallucinations with LLM judges, structured output, and context engineering: Cut errors via schema-aware chunking, dedupe and budgeted context packing, plus lightweight LLM checks and schema-constrained outputs to verify claims and enforce citations before responding.
Design effective LLM prompts for high-level control on generation output: Write prompts that steer behavior: system prompts, task decomposition, Pydantic/JSON-schema constraints, and clear rules for tone, citations, and safe refusals.
Develop end-to-end RAG applications using the software engineering best practices: Produce a maintainable codebase: clean module boundaries (ingest/retrieve/rerank/synthesize), typed configs, secrets/env management, reproducible local dev, and deploy that mirrors local.

Build Production-Ready Agentic-RAG Applications From Scratch Course: What we are going to build

Damien Benveniste — Tue, 02 Sep 2025 15:01:39 GMT

On Saturday, September 27th, I am launching a new course: Build Production-Ready Agentic-RAG Applications From Scratch! This is a fully hands-on course where we are going to deploy a production-ready Agentic-RAG application with LangGraph, FastAPI, and React! Here is what we are going to build.

Signup!

What we are going to build

We are going to build a fun web application where we can demonstrate how to orchestrate a robust RAG application using LangGraph, FastAPI, and React. Here is what we are going to build:

A user can pass a GitHub repository URL
The files of the related repository are scraped and indexed in a vector database
Now the code is available for the user to ask questions about.

On the frontend, we will need two main functionalities:

A page where we can input the repository URL and start the crawling and indexing processes:

And a chatbot interface to ask questions about the code in the repository:

On the backend, we will need the related endpoints:

The indexing endpoint will respond to the provided GitHub repository URL and the “crawl“ action to start the crawling and indexing processes.
The chat endpoint that will respond to messages sent by the user from the chatbot interface.

We are going to use the following tools:

React for the frontend
FastAPI for the backend
LangGraph for the agentic orchestration
Pinecone for the vector database
Langsmith for observability
Deploy everything on Google Cloud!

Project-based course

We will focus on building the project from the ground up, as we would on the job. Here is how we are going to structure the project development:

Introduction
- What we want to build
- Setting up the environment
The RAG Application
- The Data Parsing Pipeline
- The Indexing Pipeline
- The Basic RAG Pipeline
- Adding Observability to the Pipeline with Langsmith
- Going Agentic
The Backend Application
- The Indexing API Endpoint
- Adding Memory
- Administering the Database Data
The Frontend Application
- The Indexing Page
- The Chatbot Page
Deploying to GCP

Each session will be a live, hands-on coding session where we are going to implement every component from scratch

Going Agentic

Intent router: the entry point of the pipeline that will decide if a RAG pipeline is required.
The retriever: The sub-agent that will extract the right data
The generator: The sub-agent that will generate the response to the user

Scaling up

This is going to be a fun ride! Make sure to join us!

Signup!

New Course: Build Production-Ready Agentic-RAG Applications From Scratch

Damien Benveniste — Mon, 25 Aug 2025 15:01:14 GMT

From Prototype to Production: Ship Reliable and Scalable RAG Pipelines

The Real-World AI Engineering Roadblocks You Face Today

👋 Prototype → Production Gap — Moving from a notebook demo to a secure, observable, multi-tenant service requires orchestration, evals, guardrails, and ops most teams lack.

👋 “Easy RAG” vs “Reliable RAG” — Anyone can retrieve-then-generate; making answers faithful, fresh, fast, and cost-controlled under real traffic is the hard part.

👋 Framework Overload — The ecosystem is noisy; you need clear criteria (maturity, extensibility, latency, cost) and reference patterns to choose confidently.

👋 It’s Software Engineering First — Success hinges on clean interfaces, tests, typed configs, tracing, CI/CD, and change management—not just prompts and models.

👋 From Laptop to 1M Users — Scaling demands streaming, batching, caching, autoscaling, and SLOs, or your p95 explodes and costs spiral.

How this course will help you

✅ Ship a real Agentic RAG app, not a demo — Stand up an end-to-end stack—LangGraph → FastAPI → React, that runs locally today and deploys via a clean, fork-and-ship monorepo.

You’ll walk away with

✨ A running Agentic RAG app (LangGraph + FastAPI + React) in a fork-and-ship monorepo.

✨ An ingestion/indexing pipeline with metadata, hybrid retrieval, and optional re-ranking.

✨ A chat UI with citations, source previews, and conversation memory that behaves.

✨ Deploy scripts and env templates to go live right after class.

✨ A framework choice memo + adapters to swap models/vector stores without starting over.

Bottom line: this isn’t a vitamin, it’s a blueprint you can put in production.

What you’ll get out of this course

Orchestrate complex RAG pipelines with LangGraph and OpenAI API: Build a typed LangGraph that routes rewrite → retrieve → rerank → synthesize → cite → self-check with retries, timeouts, early-exit rules, and real tool calls, exposed as a clean HTTP API.
Build scalable asynchronous applications with FastAPI: Ship async FastAPI endpoints, well-typed request/response models, input validation, and sensible timeouts, ready to run locally and deploy to production.
Implement chatbot interfaces with React: Create a chat UI that shows citations and source previews, lets users scope queries, preserves safe chat history, and handles transient API errors gracefully.
Mitigate hallucinations with LLM judges, structured output, and context engineering: Cut errors via schema-aware chunking, dedupe and budgeted context packing, plus lightweight LLM checks and schema-constrained outputs to verify claims and enforce citations before responding.
Design effective LLM prompts for high-level control on generation output: Write prompts that steer behavior: system prompts, task decomposition, Pydantic/JSON-schema constraints, and clear rules for tone, citations, and safe refusals.
Develop end-to-end RAG applications using the software engineering best practices: Produce a maintainable codebase: clean module boundaries (ingest/retrieve/rerank/synthesize), typed configs, secrets/env management, reproducible local dev, and deploy that mirrors local.
Sign Up!

Mixture-of-Experts: Early Sparse MoE Prototypes in LLMs

Damien Benveniste — Fri, 22 Aug 2025 15:01:43 GMT

Mixture-of-Experts might be one of the most important improvements in the Transformer architecture! It allows for scaling the number of model parameters while keeping the latency associated with the forward and backward pass of the backpropagation algorithm almost constant. Scaling in the width direction, as opposed to the depth of the model, allows for keeping the gradient paths short, improving the stability of the training. We explore here 2 early models:

The Sparsely-Gated Mixture-of-Experts Layer
GShard

The First Mixture-of-Experts

The concept of the Mixture of Experts (MoE) architecture was introduced in 1991 by Jacobs et al. The idea was to combine the learning of parallel learners using a gating mechanism. The goal was to increase the model's capacity while maintaining stable training and achieving faster convergence. Deep networks tend to suffer from vanishing or exploding gradients, and extending the capacity in the width direction allows for the learning of more complex statistical patterns while keeping short gradient paths.

An "expert" E_i can be a simple feed-forward network. For example:

where W₁ⁱ and W₂ⁱ may have different dimensions depending on the expert. The gating mechanism generates a weight g_i(h) for each expert, and the MoE output is a weighted average of the experts' outputs:

g_i(h) is typically the softmax transformation of a linear projection W_g of the input features:

where W_g is a linear layer of dimension |h| ⨉ n. The softmax transformation yields a probability-like value that captures the proportion of contributions for each expert.

Early Sparse MoE Prototypes in LLMs

The Sparsely-Gated Mixture-of-Experts Layer

The Sparse mechanism

The sparse Mixture of Experts (MoE) architecture was introduced by Shazeer et al in January 2017 in LSTM-based language models as a way to drastically scale the model capacity while keeping the number of operations constant, independent of the number of experts. RNN models are hard to scale because the LSTM operations are intrinsically iterative, preventing the high parallelism provided by other computational units. Scaling in depth with more LSTM units induces high latency, while scaling with feed-forward networks limits the ability of the LSTM layers to capture long-range coherence of the input sequences. Scaling in width with many parallel experts permits the LSTM to learn the long-term dependencies while keeping the latency to a minimum.

To keep the number of operations independent of the number of experts, they introduced a routing mechanism that selects the top-k experts for each token. They tested a total number of experts that ranged between 4 and 131,072, but only the top-4 experts were used for each token.

For the sparse MoE, the architecture is the same for every expert i:

The gating mechanism is, as before, induced by a linear layer W_g mapping from hidden size d_model to n with an added normal noise ϵ:

The noise allows the model to uniformly explore the different experts in the early part of the training and prevents the collapse onto a handful of "favorite" experts. To generate the noise, the input vector is first projected with another linear layer W_n passed through a softplus transformation:

The resulting value 𝛔(h) is used as the standard deviation to generate the normal noise:

The standard deviation 𝛔(h) controls how much randomness is injected for each expert on that particular input token. Larger 𝛔(h) will lead to more exploration, and smaller 𝛔(h) will push the gate to behave almost deterministically. Softplus is basically a "soft" version of ReLU: for large positive values, it increases linearly and flattens near 0. It is used here to make sure the learned noise scale 𝛔(h) is positive and differentiable everywhere. Because the noise is input‑adaptive, the router can start off as a near‑uniform sampler (large 𝛔(h)) and gradually anneal into a confident switch (small 𝛔(h)).

Once the logits l(x) = {l₁(h), l₂(h), …, l_n(h)} have been computed, we mask the non-top-k's contribution with -∞:

And we perform a softmax transformation to obtain the contribution of each top-k expert:

The -∞ masks will lead to g_i(h) = 0 contribution for every non-top-k expert while keeping the softmax normalization as if only the top-k experts contributed to the sum. The resulting hidden state y(h) is the weighted average of top-k experts' output:

Hierarchical Mixture of Experts

In the Sparsely-Gated Mixture-of-Experts, they scaled the number of experts up to 131,072! With naive MoE, the gate and noise projection layers W_g and W_n would require a dimension (n ⨉ d_model) ~ 67M parameters (with d_model= 512), with as many operations for each token in the input sequence. At this scale, the gating mechanism becomes the bottleneck. Instead, they introduced a hierarchical gating process where the experts were grouped into a blocks of b experts. The first gate projects from d_model to a blocks:

From this first set of logits, we can pick the top-k₁ expert blocks:

Last Week to Register to the Build Production-Ready LLMs From Scratch Course!

Damien Benveniste — Wed, 09 Jul 2025 15:02:40 GMT

This Saturday, we kick off the latest cohort of the Build Production-Ready LLMs From Scratch course! This is the last week to register, so make sure to join us if you want to get the right skills as a Machine Learning engineer! This is a 6-week program to learn to build scalable LLMs from scratch and ship them to production. It will run between July 12th and Aug 17th, 2025. It includes 12 live sessions, 6 real-world hands-on projects, 64 recorded lectures, and more material.

Enroll

The Real-World LLM Engineering Roadblocks You Face Today

👋 Transitioning from General ML to LLM Specialization: You’ve built recommendation engines or classifier models, but moving into Transformer‑centric development feels like learning a whole new discipline—no clear roadmap exists.

👋 Lack of LLM‑Specific Career Path: You see “LLM Engineer” roles popping up on LinkedIn, but your current CV only shows “Data Scientist” or “ML Engineer.” You need hands‑on projects and artifacts to credibly make the jump.

👋 Career Stalled by “Academic” Skillset: You can recite Transformer papers, but when asked, “Have you shipped an LLM feature end‑to‑end?” you have no answer—and no portfolio to prove it!

👋 Prototype Meltdown Under Production Load: You’ve fine‑tuned a small model locally, but when you switch from 1 to 100 concurrent requests, your GPU memory spikes and inference grinds to a halt, because you’ve never applied continuous batching, KV caching, or paged‑attention in a live setting.

👋 RAG Integration Headaches: Turning a standalone model into a live, Retriever‑Augmented Generation service becomes a multi‑week integration nightmare.

How this course will help you

Because we’ve packaged every stage of the LLM lifecycle, from career transition to production rollout, into a six‑week bootcamp that:

✅ Guides Your Career Pivot: You’ll emerge with six polished GitHub projects, a deployment playbook, and RAG demos that transform your resume from “ML generalist” to “LLM Specialist.”

✅ Attacks Each Pain‑Point Head‑On: Attacks each pain point head‑on with six job‑mirroring projects (from scratch → RLHF → scaling → deployment → RAG), so you never waste time on dead‑end tutorials

✅ Live Code‑Along Workshops & Office Hours: Tackle your own fine‑tuning bugs, scaling hiccups, and deployment errors alongside Damien in dedicated sessions, so you get hands‑on fixes for the exact issues you’ll face on the job.

✅ Ready‑to‑Use Repos & Playbooks: Grab our curated starter code, development scripts, deployment templates, and debugging checklists, so you can plug them straight into your next project without reinventing the wheel.

✅ A Portfolio of Six Production‑Grade Projects: Leave with six end‑to‑end deliverables, from a Transformer built from scratch to a live RAG API, ready to showcase on GitHub, in performance reviews, or to hiring managers.

No more scattered blog-hopping or generic bootcamps, this is the only cohort where you’ll master Transformer internals and ship production‑grade LLM systems while making the career leap you’ve been aiming for.

What You’ll Actually Build and Ship

Across six hands‑on projects, you’ll deliver deployable LLM components and applications, no fluff, just job‑ready code:

✅ A Modern Transformer Architecture from scratch: Implement a sliding‑window multihead attention to slash O(N²) to O(N·w), RoPE for relative positional encoding, and the Mixture-of-Expert architecture for improved performance, all in PyTorch.

✅ Instruction‑Tuned LLM: Fine‑tune a model with supervised learning, RLHF, DPO, and ORPO for instruction following on a real benchmark and compare performance gains.

✅ Scalable Training Pipeline: Containerize a multi‑GPU job with DeepSpeed ZeRO on SageMaker to maximize throughput and minimize cost.

✅ Extended‑Context Model: Modify RoPE scaling, apply 4/8‑bit quantization, and inject LoRA adapters to double your context window.

✅ Multi‑Mode Deployment: Stand up a Hugging Face endpoint, a vLLM streaming API, and an OpenAI‑compatible server, all Dockerized and optimized for low latency.

✅ End‑to‑End RAG Chat App: Build a FastAPI backend with conversational memory and a Streamlit UI for live Retrieval‑Augmented Generation.

By the end of Week 6, you won’t just know these techniques, you’ll have shipped six production‑grade artifacts, each reflecting the exact pipelines, optimizations, and deployment routines you’ll use on the job.

Live & Recorded Content: Reinforce, Deepen, Accelerate

✨ 12 Interactive Live Workshops (3 hrs each): Each session follows the Concept → Code flow. I’ll introduce the day’s core topic (e.g. self-attention, LoRA, vLLM optimizations, ...), and we’ll implement the features step‑by‑step in code so you see exactly how theory maps to code. Bring your questions!

✨ 10 + Hours of On‑Demand Deep‑Dive Lectures: Short videos (10–20 min) on Transformer internals, fine-tuning tricks, deployment optimizations. Watch before each project to hit the ground running. Step through every line of code at your own pace; perfect for review or catching up if you miss a live session. Downloadable slide decks, annotated notebooks, and cheat sheets you’ll reference long after graduation.

Why This Matters: Live workshops turn recorded concepts into actionable skills. You’ll see how theory maps directly onto code, get instant feedback, and internalize best practices. Then, recorded lectures become your asynchronous safety net, letting you revisit tricky topics, prepare for upcoming labs, and solidify your understanding on demand.

Let me know if you have any questions. I hope to see you there!

Build Production-Ready LLMs From Scratch Starting on July 12th!

Damien Benveniste — Mon, 16 Jun 2025 15:02:40 GMT

Get ready! The latest iteration of the Build Production-Ready LLMs From Scratch live course is starting on July 12th! This is a 6-week program to learn to build scalable LLMs from scratch and ship them to production. It will run between July 12th and August 17, 2025. It includes 12 live sessions, 6 real-world hands-on projects, 64 recorded lectures, and more material. The first 30 people to sign up will get a 20% discount by applying the promo code FIRST! So make sure to sign up early:

Signup

The Real-World LLM Engineering Roadblocks You Face Today

👋 RAG Integration Headaches: Turning a standalone model into a live, Retriever‑Augmented Generation service becomes a multi‑week integration nightmare.

How this course will help you

Because we’ve packaged every stage of the LLM lifecycle, from career transition to production rollout, into a six‑week bootcamp that:

✅ Guides Your Career Pivot: You’ll emerge with six polished GitHub projects, a deployment playbook, and RAG demos that transform your resume from “ML generalist” to “LLM Specialist.”

What You’ll Actually Build and Ship

Across six hands‑on projects, you’ll deliver deployable LLM components and applications, no fluff, just job‑ready code:

✅ Instruction‑Tuned LLM: Fine‑tune a model with supervised learning, RLHF, DPO, and ORPO for instruction following on a real benchmark and compare performance gains.

✅ Scalable Training Pipeline: Containerize a multi‑GPU job with DeepSpeed ZeRO on SageMaker to maximize throughput and minimize cost.

✅ Extended‑Context Model: Modify RoPE scaling, apply 4/8‑bit quantization, and inject LoRA adapters to double your context window.

✅ Multi‑Mode Deployment: Stand up a Hugging Face endpoint, a vLLM streaming API, and an OpenAI‑compatible server, all Dockerized and optimized for low latency.

✅ End‑to‑End RAG Chat App: Build a FastAPI backend with conversational memory and a Streamlit UI for live Retrieval‑Augmented Generation.

Live & Recorded Content: Reinforce, Deepen, Accelerate

Let me know if you have any questions. I hope to see you there!

Last Week to Register to the Build Production-Ready LLMs From Scratch Course!

Damien Benveniste — Mon, 19 May 2025 15:54:44 GMT

This Saturday, we kick off the Build Production-Ready LLMs From Scratch course! This is the last week to register, so make sure to join us if you want to get the right skills as a Machine Learning engineer! This is a 6-week program to learn to build scalable LLMs from scratch and ship them to production. It will run between May 24th and June 29, 2025. It includes 12 live sessions, 6 real-world hands-on projects, 64 recorded lectures, and more material.

Enroll

The Real-World LLM Engineering Roadblocks You Face Today

👋 RAG Integration Headaches: Turning a standalone model into a live, Retriever‑Augmented Generation service becomes a multi‑week integration nightmare.

How this course will help you

Because we’ve packaged every stage of the LLM lifecycle, from career transition to production rollout, into a six‑week bootcamp that:

✅ Guides Your Career Pivot: You’ll emerge with six polished GitHub projects, a deployment playbook, and RAG demos that transform your resume from “ML generalist” to “LLM Specialist.”

What You’ll Actually Build and Ship

Across six hands‑on projects, you’ll deliver deployable LLM components and applications, no fluff, just job‑ready code:

✅ Instruction‑Tuned LLM: Fine‑tune a model with supervised learning, RLHF, DPO, and ORPO for instruction following on a real benchmark and compare performance gains.

✅ Scalable Training Pipeline: Containerize a multi‑GPU job with DeepSpeed ZeRO on SageMaker to maximize throughput and minimize cost.

✅ Extended‑Context Model: Modify RoPE scaling, apply 4/8‑bit quantization, and inject LoRA adapters to double your context window.

✅ Multi‑Mode Deployment: Stand up a Hugging Face endpoint, a vLLM streaming API, and an OpenAI‑compatible server, all Dockerized and optimized for low latency.

✅ End‑to‑End RAG Chat App: Build a FastAPI backend with conversational memory and a Streamlit UI for live Retrieval‑Augmented Generation.

Live & Recorded Content: Reinforce, Deepen, Accelerate

Let me know if you have any questions. I hope to see you there!

All About The Modern Positional Encodings In LLMs

Damien Benveniste — Mon, 28 Apr 2025 15:02:32 GMT

The Positional Encoding in LLMs may appear somewhat mysterious the first time we come across the concept, and for good reasons! Over the years, researchers have found many different ad-hoc ways to encode positions and relative positions in the attention mechanism. The most recent surprise to me is the emergence of NoPE (no positional encoding) as being better for generalization on longer sequences. Today, we look at:

Additive Relative Positional Embeddings
Multiplicative Relative Positional Embeddings
ALiBi: Attention With Linear Biases
RoPE: Rotary Position Embedding
- The Complex Number Representation
- Increasing The Context Window With RoPE
No Position Encoder: NoPE
- How NoPE Can Learn Relation Positional Information
- Better Generalization for Longer Distance
- Llama 4's iRoPE

The original positional encoding scheme introduced in "Attention is All You Need" uses sinusoidal functions to create absolute position representations, but this approach revealed significant shortcomings when handling longer sequences or transferring to different sequence lengths:

Fixed context windows and poor extrapolation: When faced with sequences longer than this predefined limit, the model must either truncate the input or use positional values it never encountered during training. Perhaps the most significant limitation is that vanilla encodings represent absolute positions rather than relationships between tokens.
The absolute vs. relative position problem: In language, relative positions often matter more than absolute ones. Consider the sentence: "The dog that chased the cat ran away." The relationship between "dog" and "ran" remains the same whether this is the opening sentence of a document or appears on page fifty, but absolute encodings fail to directly capture this invariance.
Information dilution through layers: As signals propagate through the multiple layers of a Transformer, the influence of the original positional information tends to weaken. The model must work harder to maintain positional awareness in deeper layers, especially for distant tokens.
Limited inductive bias for local relationships: Natural language exhibits a locality bias as nearby words often have stronger relationships than distant ones. The vanilla encoding doesn't naturally encode this bias, treating position 5 and position 500 as equally valid attention targets from a structural perspective, requiring the model to learn this pattern from data alone.
Mathematical constraints: The sinusoidal functions used in the original formulation were chosen partly for their theoretical ability to generalize to unseen positions. However, in practice, they still struggle with significant extrapolation beyond the training range. The model learns to associate specific patterns with specific position ranges, and these associations become increasingly unreliable as we move farther from the training distribution.

These limitations collectively motivated researchers to develop more sophisticated positional encoding schemes, from Shaw's direct relative encodings to Transformer-XL's decomposed approach, ALiBi's distance-based penalties, and RoPE's elegant rotational solution, each addressing different aspects of these fundamental challenges.

You can find the discussion about the original positional encoding here:

Additive Relative Positional Embeddings

One of the first works addressing capturing the relative position between tokens instead of the absolute ones was done by Shaw et al in 2018. It addressed many of the original positional encoding's shortcomings. Instead of having an embedding matrix added to the semantic representation of the tokens, it modifies directly the computation of the attention layers with two positional embeddings. The first one a_ij^K is added to the key representation, and the second a_ij^V, to the value representation:

where a_ij^K and a_ij^V are learned embeddings of size d_model that depend only on the relative distance j-i. Those embeddings are learned in each attention layer and influence more directly the interactions between tokens. Ignoring the attention heads for simplicity, the computation of a context vector c_i is modified as:

They clip the maximum relative distance to a constant L, learning only 2L+1 position embeddings for each positional encoding. If we call the matrices of learned positions w^K and w^V, the actual relative position representations that are added to keys and values are the clipped version of those matrices:

where

which clamps the relative distance to the range [−L, L]. This means that for any relative distance j-i > L,

and any relative distance j-i < -L,

In practice, L can be chosen quite small as the performance remained relatively stable for L > 2 (testing L = 4, L = 16, L = 64, and L = 256). This suggests that distinguishing the precise relative distances of tokens becomes less important beyond a few positions away. The model primarily needs fine-grained position information for nearby contexts, while more distant relationships can be effectively captured with less positional precision, allowing the content representations themselves to drive attention patterns at longer distances. Even with a very small clipping window (L = 2), the model with relative positional encoding achieved substantial improvements over absolute positioning.

Despite the added performance gain, it is important to note that it is at the cost of added parameters and added time complexity. At the time, Shaw et al. performed their experiments with a 65M parameters transformer model with 6 layers and d_model = 512. We need two new parameter layers of size (2L+1) x d_model, and with L = 16, it is ~200K additional parameters, which is negligible compared to the overall size of the model. More importantly, computing all the relative distances between the N tokens, the constant memory access, and adding the related embeddings is an O(N²) process that imposes a significant additional computational work on top of the already quadratic complexity of attention.

Multiplicative Relative Positional Embeddings

In a previous newsletter, we introduced the attention mechanism developed with Transformer-XL to handle arbitrarily long sequences, but we left out the discussion about the positional encoding. Developing a new relative positional encoding was one of the critical pieces to handle longer sequences. The approach is based on the following analysis we provided in a previous newsletter for the original positional encoding:

In this equation, we decomposed the contributions from the different content and positional components for the first attention layer in the model when using the vanilla absolute positional encoding. We showed how R(k) was capturing the relative positional information between tokens, and we hope for the model to learn W^Q and W^K such that it can effectively utilize that information. To help the model capture better the content and position interactions between tokens, the relative positional encoding introduced in Transformer-XL modified this original equation. The relative positional encoding will now be applied within every layer following this strategy:

x_i^T W^QT W^K x_i+k originally captured the pure content-based interaction between tokens. x_i, in the original Transformer, represents the embedding vector from the token embedding, and the related hidden state h_i is the sum of the token embedding and the positional encoding:
We modify this interaction by applying it directly to the hidden states h_i and h_i+k where k is the distance between the related tokens:
We specifically distinguish the projection W_E^K to be content-specific.
x_i^T W^QT W^K PE(i) R(k) or equivalently x_i^T W^QT W^K PE(i+k) R(k) captured a token's preference for attending to positions was tied to fixed locations. However, this is confusing for the model because PE(i+k) depends on the specific position of x_i. Instead, in the Transformer-XL, they modified this interaction only accounting for the relative distance k from the hidden state h_i:
r_k is a vector of the positional encoding matrix R. As for the additive relative positional embeddings, R is a (2L+1) x d_model matrix where r_k only depends on the relative distance between h_i and h_i+k, and L is a constant chosen to clip the maximum distances that can be represented. However, R is not a learned parameter matrix, but a fixed encoding as in the original Transformer architecture. W_R^K (different from W_E^K) is a specialized weight matrix that further enhances this separation of positional processing.
PE(i)^T W^QT W^K x_i+k is a position-dependent bias for attending to content. This is a strange term because it has more or less weight depending on the position of the query. To make it position independent, we replace PE(i)^T W^QT with a global learnable parameter u^T:
u is a vector of size d_model. This reflects an insight that the attentive bias toward different content should remain consistent regardless of the query position. In other words, the importance of certain word types doesn't need to depend on position, so a single global parameter can replace position-specific queries. This simplifies the model while maintaining expressiveness.
PE(i)^T W^QT W^K PE(i+k) is a position-position iteration term that depends on the specific position of the query. To make it position independent while keeping the relative positional information, we introduce another global learnable parameter v^T of size d_model to replace PE(i)^T W^QT, and use again the relative positional encoding R:
As before, we use the position-specific projection W_R^K. The intuition is that certain relative distances might be generally more important than others, regardless of the absolute positions involved. For example, adjacent tokens (small k) might generally be more related than distant ones, regardless of their absolute positions in the sequence.

This leads to redefining the alignment scores as:

with W^Qh_i = q_i, the query and W_E^Kh_i+k = k_i+k, the key, we can regroup the terms:

with

being a pure content interaction term and

representing a content-to-relative distance interaction term. W_E^Kh_i+k = k_i+k can be thought of as a content contribution to the key, while W_R^Kr_k is the relative positional piece of the key. The global parameters u and v can be seen as providing "default query projections" that are active regardless of the specific query token. This ensures that important content and positional patterns always receive some attention.

R follows the same sinusoidal encoding function as the original Transformer, just applied to relative positions instead of absolute positions:

where k is the relative distance between positions, ranging from −L to +L, and j is the dimension index from 0 to d_model.

Computing

requires careful thought to limit the complexity of the problem. With N tokens, we could naively extract all the N² related r_k since there are N² pairs of tokens. However, there would be many duplicate vectors since different pairs of tokens have the same distance between them. Instead, we can directly use the whole R matrix and pass it through the linear layer W_R^K:

R is a (2L+1) x d_model, and W_R^K is a d_model x d_model matrix, therefore the "positional keys" K_R is a (2L+1) x d_model matrix. As a consequence

is a (2L+1) dimensional vector which leads to a (2L+1) x N content to position alignment matrix for N queries.

Let's make sure we understand the different tensors in play. We have the queries Q = [q₁, …, q_N], the keys K = [k₁, …, k_N], and the positional equivalent to the keys K_R =[W_R^Kr_-L, …, W_R^Kr_L]. The relative positions in K_R are not aligned with the relative positions of K with respect to Q. This means we cannot simply add

because of the misalignment in the key ordering. We just need to change the indices of the elements in (Q + v)^T K_R. For query position i, we need:

For key position 0: a relative distance i−0 = i
For key position 1: a relative distance i−1
For key position 2: a relative distance i−2
And so on...

In other words, for the i-th row of the content to position alignment matrix P = (Q + v)^T K_R, we need to select values in a specific pattern. The solution is to reshape and shift the P matrix so that the correct relative position scores align with the positions we need in the final attention matrix.

In the case of Transformer-XL, remember that the sequences are processed in segments of size n « N, and the attention mechanism per segment follows an O(n²) time complexity. W_R^KR also operates on segments, and the time complexity is

The shift operation needs at least O(n²) operations to reindex the alignment score matrix. Therefore, the overall time complexity associated with the relative positional encoding follows the asymptotic behavior O(n²).

Transformer-XL showed improvements in the model's ability to utilize longer contexts. When trained with segments of length 128 but evaluated with various attention lengths, the model showed continued improvements in perplexity up to 640 tokens. The paper reported improved perplexity metric when increasing evaluation context length beyond training length, something absolute encoding couldn't achieve.

ALiBi: Attention With Linear Biases

The relative positional encoding developed in Transformer-XL is designed to handle very long sequences, but it adds complexity to the attention layer, and, while it can handle longer contexts through its recurrence mechanism, it was not specifically designed for extrapolation to arbitrary lengths beyond training. ALiBi (Attention With Linear Biases) was introduced in 2021 as a simpler approach that could easily be extrapolated to much longer sequences than the ones seen during training. It does not require any model parameters and can be expressed as a penalty on the alignment score:

where j-i is the distance between the query q_i and the key k_j, and m_h is a head-specific constant, with h being the index of the head. In the case of causal language modeling, we always have j-i ≤ 0, leading to lower attention for far-away tokens. ALiBi tends to sacrifice true long-range modeling for extrapolation capability. m_h is chosen as:

where h ∈ [0, …, n_head -1]. For example, if we have 8 heads, the first head would have m₀= 2⁰= 1, and the last head m₇ = 2^-7 = 0.0078125, with a spectrum of the different intermediary slopes in between. For j-i = 1000, for example, we get m₇(j-i)~ -7.8, which is a substantial penalty to capture long-range interaction between tokens. This suggests that ALiBi may not be optimal for tasks requiring very long-range dependencies (like book-length coherence), but it represents a valuable engineering trade-off that improves efficiency without sacrificing performance on many practical tasks.

Despite this limitation, ALiBi works well for extrapolation for several reasons:

Graduated attention ranges: The different head slopes create a spectrum of attention distances. While no head truly specializes in very long-range attention, the collection of heads creates a gradient of focus distances.
Local coherence dominance: Language has a hierarchical structure where local coherence (within paragraphs or nearby sentences) often matters more than very distant relationships. The bias aligns with this natural property of language.
Information propagation: Information can still flow across long distances through multiple layers of the transformer. Even if direct attention across 1000 tokens is penalized, information can propagate through intermediate positions across layers.
Relative vs. absolute positioning: Unlike sinusoidal embeddings that break down completely outside their training range, ALiBi's linear bias at least provides a consistent, predictable signal at any distance.

The paper showed that a model trained on 512 tokens could handle sequences of 3072 tokens with better perplexity than a sinusoidal model trained on 3072 tokens. A 1.3 billion parameter model trained on 1024 tokens achieved the same perplexity as a sinusoidal model trained on 2048 tokens when evaluated on 2048-token sequences. As input length increases beyond training length, sinusoidal models' performance degrades almost immediately while ALiBi's performance continues improving up to ~2-3x training length before plateauing.

RoPE: Rotary Position Embedding

The Rotary Position Embedding (RoPE) is now one of the most common strategies used to inject the relative positional information within the attention mechanism. The idea behind RoPE is to rotate the keys and queries based on the position of the related tokens in the input sequences. This will inject the absolute positional information directly into the queries and keys. Let's look at a toy example to understand the logic. Let's consider a 2-dimensional query q_i and a 2-dimensional key k_j. To rotate 2-dimensional vectors, we use rotation matrices:

Join us for a Free LIVE Coding Event: Build The Self-Attention in PyTorch From Scratch

Damien Benveniste — Fri, 25 Apr 2025 15:00:39 GMT

Next Friday, I am inviting you to join me for an exciting live coding event. It is a completely free event where I will explain the basics of the self-attention layer and implement it from scratch in PyTorch. From the vanilla self-attention to the multi-head attention layer, I will walk you through all the little details.

Over the past few years, I have realized that the logic behind self-attention still eludes many people who want to dive deeper into the field. For me, implementing from scratch is the best way to learn this core element of every LLM. Once the self-attention implementation becomes more intuitive, it opens the doors to understanding all the small improvements that have led to the level of maturity we have today in the field of LLMs.

Signup

This is an opportunity to talk face-to-face and ask questions. The event is on May 2nd at 9:30 AM PST. I hope to see you there!

Build Production-Ready LLMs From Scratch

Damien Benveniste — Mon, 21 Apr 2025 15:03:03 GMT

Big news! I am now partnering with Maven as an instructor to teach the Build Production-Ready LLMs From Scratch live course! This is a 6-week program to learn to build scalable LLMs from scratch and ship them to production. It will run between May 24th and June 29, 2025. It includes 12 live sessions, 6 real-world hands-on projects, 64 recorded lectures, and more material. The first 30 people to sign up will get a 20% discount by applying the promo code FIRST! So make sure to sign up early:

Signup

The Real-World LLM Engineering Roadblocks You Face Today

👋 RAG Integration Headaches: Turning a standalone model into a live, Retriever‑Augmented Generation service becomes a multi‑week integration nightmare.

How this course will help you

Because we’ve packaged every stage of the LLM lifecycle, from career transition to production rollout, into a six‑week bootcamp that:

✅ Guides Your Career Pivot: You’ll emerge with six polished GitHub projects, a deployment playbook, and RAG demos that transform your resume from “ML generalist” to “LLM Specialist.”

What You’ll Actually Build and Ship

Across six hands‑on projects, you’ll deliver deployable LLM components and applications, no fluff, just job‑ready code:

✅ Instruction‑Tuned LLM: Fine‑tune a model with supervised learning, RLHF, DPO, and ORPO for instruction following on a real benchmark and compare performance gains.

✅ Scalable Training Pipeline: Containerize a multi‑GPU job with DeepSpeed ZeRO on SageMaker to maximize throughput and minimize cost.

✅ Extended‑Context Model: Modify RoPE scaling, apply 4/8‑bit quantization, and inject LoRA adapters to double your context window.

✅ Multi‑Mode Deployment: Stand up a Hugging Face endpoint, a vLLM streaming API, and an OpenAI‑compatible server, all Dockerized and optimized for low latency.

✅ End‑to‑End RAG Chat App: Build a FastAPI backend with conversational memory and a Streamlit UI for live Retrieval‑Augmented Generation.

Live & Recorded Content: Reinforce, Deepen, Accelerate

Let me know if you have any questions. I hope to see you there!

Chapter 4 of The Big Book of Large Language Models is Here!

Damien Benveniste — Mon, 31 Mar 2025 15:02:52 GMT

Chapter 4 of the Big Book of Large Language Models is finally here!

That was a difficult chapter to write! Originally, I wanted to cram in that chapter all the improvements related to the Transformer architecture, since the Attention is all you need paper, but I realized that it would be too long for one chapter. I ended up focusing only on improvements related to the attention layer and delaying things like relative positional encoding and Mixture of Experts to the next chapter. In this chapter, I addressed the following improvements:

Sparse Attention Mechanisms
- The First Sparse Attention: Sparse Transformers
- Choosing Sparsity Efficiently: Reformer
- Local vs Global Attention: Longformer and BigBird
Linear Attention Mechanisms
- Low-Rank Projection of Attention Matrices: Linformer
- Recurrent Attention Equivalence: The Linear Transformer
- Kernel Approximation: Performers
Memory Efficient Attention
- Self-attention Does Not Need O(N 2) Memory
- The FlashAttention
Faster Decoding Attention Mechanisms
- Multi‑Query Attention
- Grouped‑Query Attention
- Multi-Head Latent Attention
Long Sequence Attentions
- Transformer-XL
- Memorizing Transformers
- Infini-Attention

Obviously, I could not include everything that was ever invented in the context of the attention layer, but I believe those use cases capture well the different research routes that have been explored since then. I believe it is a very important chapter, as most materials available online tend to focus on the vanilla self-attention, which starts to be an outdated concept for today’s standards. I also found that trying to understand how to improve the self-attention is a very good way to understand what it is we are trying to improve in the first place! The self-attention may appear odd at first, but diving into the inner workings of the layer in order to improve it gives us a level of understanding that is beyond anything we can learn just by looking at the original self-attention. I hope you will enjoy it!

Reduce AI Model Operational Costs With Quantization Techniques

Damien Benveniste — Wed, 26 Mar 2025 15:01:25 GMT

Model quantization is becoming a core strategy for training and deployment! I am excited to introduce you to Louis-François Bouchard! He is an exceptional AI educator and entrepreneur, and in this guest post, he presents the fundamentals for model quantization and a detailed tutorial on how to quantize a Llama 3 model.

Louis-François Bouchard, a dedicated AI educator and entrepreneur since 2019, left his PhD studies after recognizing a disconnect between academic research and industry needs. As the founder of Towards AI, he is committed to making artificial intelligence accessible and bridging that gap through practical teaching on the Towards AI Academy platform. With a wealth of free resources online—videos, blogs, and newsletters—his academy empowers a diverse global community of developers and enthusiasts to innovate and thrive with new relevant AI technologies like LLMs and everything around them.

Large AI models are changing industries worldwide, yet their enormous size makes them challenging to deploy efficiently. With billions of parameters, they demand powerful GPUs, abundant VRAM, and extensive compute resources, leading to high memory usage and steep operational costs.

Model quantization has emerged as a powerful technique to address these issues. By reducing the precision of a model’s weights, quantization dramatically cuts memory footprints. A quantized model can often run significantly faster and use a fraction of the memory of its full-precision equivalent, lowering inference latencies and hardware costs with minimal impact on accuracy.

In this article, we’ll explore the fundamentals of model quantization, examining its underlying principles, various precision levels, and its practical implementation through a detailed code example with the Hugging Face bitsandbytes library. We will guide you step-by-step on how to load a full-precision Meta Llama 3 model, convert it into a 4-bit quantized version, and compare their memory usage, inference speed, and output quality. Additionally, we will explore the trade-offs and best practices needed to optimize model performance while achieving significant memory savings and faster inference.

What is Model Quantization?

Quantization is a method for shrinking neural network models, including Transformers, by reducing the precision of their parameters (weights, biases) and activations. Lower precision reduces a model’s memory footprint and computational requirements, enabling deployment on resource-constrained devices like mobile phones, smartwatches, and embedded systems.

A model like Meta Llama3 8B, which contains 8 billion parameters, stores these parameters in model weight files loaded onto GPUs for inference. These weights are essentially matrices stored in different numerical precisions. By quantizing these weights (reducing precision), you decrease the GPU compute and memory requirements. However, overly aggressive quantization can sometimes reduce inference accuracy.

Many open-source LLMs accessed through cloud APIs or downloaded locally are already quantized. Providers typically convert models from higher precision (FP32 or FP16) to lower precision formats (INT8 or 4-bit) to optimize performance. Properly executed quantization can significantly reduce hosting and deployment costs while preserving most of the original accuracy.

Think of this: when someone asks you the time, saying “about 11 p.m.” is faster but less precise than “10:58 p.m.” This is how quantization works. It accelerates processing at the expense of slight accuracy losses. The exact trade-off depends on the numeric format chosen (FP16, BFLOAT16, INT8, etc.).

Floating-point precision determines how accurately data is stored and processed in machine learning. Higher precision (e.g., Float32) offers better accuracy but requires more memory, whereas lower precision types (Float16, BFloat16) reduce memory usage at the cost of some precision. The figure below illustrates how different floating-point formats allocate bits to sign, range, and precision:

For instance, a Meta Llama2 70B model using FP16 precision consumes roughly 130 GB:

(70,000,000,000 × 2 bytes) / 1024³ ≈ 130.385 GB

Further quantization to 8-bit or 4-bit reduces memory and storage even more. Different inference providers (e.g., Together.ai or Groq) use varying quantization schemes, affecting performance across identical models.

At its core, quantization is simple, it’s okay if you didn’t fully grasp the equation above. You just need to remember that quantization trades a little precision for efficiency, we discuss this trade-off in detail in the next section. For LLMs running at scale, this trade-off doesn’t exist in isolation. You’ll need to weigh it alongside factors like inference speed, memory efficiency, and deployment constraints. If you’re want to learn more about how these techniques work in a practical environment, you might find our From Beginner to Advanced LLM Developer course quite useful. We discuss this trade-off in detail in the next section.

Precision Levels and Memory Savings

Quantization can use various numerical precisions. Each precision level has a different trade-off in terms of memory, computational speed, and model accuracy. The common ones for AI models are FP32, FP16 (or BF16), INT8, INT4, and, more recently, a special 4-bit float (NF4). The table below summarizes these precision levels, their memory costs relative to 32-bit precision (FP32), and their key characteristics:

For example, FP16 cuts memory usage in half with minimal performance impact, often speeding up inference on GPUs that support half-precision. Similarly, storing parameters in 8-bit integers (INT8) rather than 32-bit floating-point (FP32) significantly decreases the model size and computational load, reducing the model size by roughly four times when calibrated to minimize accuracy loss, making it attractive for deployment.

While INT4 offers even greater compression, potentially up to eight times smaller, the effective reduction is closer to six times due to overhead and the retention of some values in higher precision to maintain accuracy. NF4 addresses this challenge by preserving more information than INT4, proving especially useful for fine-tuning LLMs using methods like QLoRA.

Using fewer bits not only reduces memory usage but also enhances computational speed by enabling faster data transfers and the use of specialized low-precision instructions.

Types of Quantization

Not all quantization approaches are the same; there are multiple strategies to quantize a model, each with its procedure and use case. The two broad categories are Post-Training Quantization (PTQ) and Quantization-Aware Training (QAT).

Post-Training Quantization (PTQ): PTQ is applied after a model is trained. The weights of a pre-trained model are converted to a lower precision in a single calibration step without further training. This approach is fast, does not require the full training dataset, and uses a small calibration dataset (a few hundred samples) to estimate value ranges for quantization. The process quantizes the weights (and optionally activations), dramatically reducing the model size in minutes. The main drawback is a potential small accuracy drop if the quantization error isn’t fully addressed. PTQ is an excellent choice when quick optimization is needed or retraining resources are limited. Modern PTQ methods, such as GPTQ and AWQ for LLMs, can reduce weights to 4 bits with minimal accuracy loss. Key variants include:

Weight-only Quantization:
This strategy compresses the model’s weights, often the largest memory component, while keeping activations (inputs/outputs of each layer) in higher precision (FP16/FP32) to avoid additional errors. This approach reduces model memory size and can partially speed up inference by using lower precision for weight matrices. Many recent LLM quantization methods focus on weight-only quantization to preserve accuracy, offering substantial memory savings given that LLM weights can be hundreds of GB at FP32.
Full Quantization (Weights + Activations):
Full quantization compresses both weights and activations, often using int8 or int16 formats for additional speedups. However, quantizing activations can be challenging for LLMs due to outlier values, channels with very large magnitude activations that, when quantized, may introduce significant errors. Naively quantizing these outliers can either hurt accuracy through clipping or lead to underflow. A common workaround is mixed precision, which retains higher precision for outlier activations while quantizing the rest. For instance, LLM.int8 detects outlier features in each layer and processes them in 16-bit while handling most multiplications in 8-bit.

Quantization-Aware Training (QAT): QAT ****involves quantizing model weights during training or fine-tuning. Weights are rounded to lower precision (like 8-bit) for calculations but stored and updated in higher precision (32-bit). This process is called “fake quantization.” This process allows the model to adapt its parameters to compensate for rounding errors, resulting in higher accuracy at a given bit-width. However, QAT requires additional training time and data, making it impractical for very large models. It is often used on smaller models or when maximum accuracy is essential, while smart PTQ methods are generally preferred for large language model deployments where retraining is not feasible.

Now that we’ve covered the “what” and “why” of quantization, let’s dive into the “how.” In the next section, we look at the different techniques to perform quantization.

Quantization Techniques with Code Examples

There are different techniques to perform quantization, from straightforward uniform quantization of each weight (scalar quantization) to more complex methods tailored for LLMs. In this section, we’ll explore a few key methods using code:

Scalar Quantization

Scalar quantization treats each dataset dimension independently. First, it calculates the minimum and maximum values for each dimension and then segments the range into uniform intervals (bins). Each value is assigned to a bin, effectively quantizing the data.

For example, let’s execute scalar quantization on a dataset with 2000 vectors (each 256-dimensional) generated from a Gaussian distribution:

import numpy as np

dataset = np.random.normal(size=(2000, 256))

# Calculate and store minimum and maximum across each dimension
ranges = np.vstack((np.min(dataset, axis=0), np.max(dataset, axis=0)))

Next, we determine the start and step for each dimension. Here, we use 8-bit unsigned integers (uint8), which provide 256 bins:

starts = ranges[0,:]
steps = (ranges[1,:] - ranges[0,:]) / 255

The quantized dataset is calculated as follows:

scalar_quantized_dataset = np.uint8((dataset - starts) / steps)

The scalar quantization process can be encapsulated in a function as below:

def scalar_quantisation(dataset):
    # Calculate and store minimum and maximum across each dimension
    ranges = np.vstack((
        np.min(dataset, axis=0), 
        np.max(dataset, axis=0)
    ))
    starts = ranges[0,:]
    steps = (ranges[1,:] - starts) / 255
    return np.uint8((dataset - starts) / steps)

Product Quantization

While scalar quantization treats each dimension independently, it may not account for the data distribution, potentially causing significant information loss. Consider the following vectors:

array = [
	[8.2, 10.3, 290.1, 278.1, 310.3, 299.9, 308.7, 289.7, 300.1],
	[0.1, 7.3, 8.9, 9.7, 6.9, 9.55, 8.1, 8.5, 8.99]
]

Applying scalar quantization to convert these vectors to a 4-bit integer leads to a considerable loss of information:

quantized_array = [
	[0, 0, 14, 13, 15, 14, 14, 14, 14]
	[0, 0, 0, 0, 0, 0, 0, 0, 0]
]

Product quantization enhances this approach by splitting the original vector into sub-vectors and quantizing each of these sub-vectors separately. With product quantization, you can:

Split each vector in the dataset into m separate sub-vectors.
Group the data in each sub-vector into k centroids, utilizing techniques such as k-means clustering.
Substitute each sub-vector with the index of the closest centroid from the relevant codebook.

For example, with m = 3 sub-vectors and k = 2 centroids:

from sklearn.cluster import KMeans
import numpy as np

# Given array
array = np.array([
    [8.2, 10.3, 290.1, 278.1, 310.3, 299.9, 308.7, 289.7, 300.1],
    [0.1, 7.3, 8.9, 9.7, 6.9, 9.55, 8.1, 8.5, 8.99]
])

# Number of subvectors and centroids
m, k = 3, 2

# Divide each vector into m disjoint sub-vectors
subvectors = array.reshape(-1, m)

# Perform k-means on each sub-vector independently
kmeans = KMeans(n_clusters=k, random_state=0).fit(subvectors)

# Replace each sub-vector with the index of the nearest centroid
labels = kmeans.labels_

# Reshape labels to match the shape of the original array
quantized_array = labels.reshape(array.shape[0], -1)

# Output the quantized array
quantized_array

# Result
> array([[0, 1, 1],
       [0, 0, 0]], dtype=int32)

By storing only the centroid indices, product quantization reduces memory usage and can speed up nearest-neighbor searches. It balances memory footprint and accuracy, depending on the number of centroids and sub-vectors used.

LLM-Specific Quantization Methods (GPTQ, AWQ, LLM.int8)

More advanced quantization techniques have been developed to address the challenges of maintaining accuracy in LLMs while effectively reducing their size. Let’s look at a few notable ones:

LLM.int8(): This technique identifies that activation outliers (significantly different values) disrupt the quantization of larger models. The proposed solution is to retain these outliers in higher precision, thus ensuring the model’s performance is not adversely affected.
GPTQ: GPTQ (post-training quantization for GPT models) quantizes each layer individually, minimizing the mean squared error (MSE) between quantized and full-precision weights. It uses a mixed int4-fp16 scheme, quantizing weights as int4 while keeping activations in float16, with real-time de-quantization during inference (GPTQ paper).
AWQ: identifies a small percentage (0.1%-1%) of critical weights based on activation magnitude and avoids quantizing them, preserving vital information in FP16 format. This technique balances efficiency with performance, though it introduces mixed-precision data types that may require additional scaling to ensure uniformity (AWQ paper).

Image from the “AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration” paper

QLoRA (Quantization + Low-Rank Adaptation)

QLoRA uses quantization for efficient fine-tuning of LLMs. By quantizing a pre-trained model to 4-bit and then training small Low-Rank Adaptation (LoRA) matrices on top, QLoRA makes fine-tuning accessible even for large models (up to 65B parameters) on a single GPU. This approach employs the 4-bit NormalFloat (NF4) data type, optimized for weights following a normal distribution. Quantile quantization ensures each bin contains an equal number of values, minimizing quantization error. With standardized weights via σ scaling, QLoRA matches full 16-bit fine-tuning performance on NLP tasks (QLoRA paper) and has been highlighted in industry blogs for its efficiency (Hugging Face blog).

From “QLoRA: Efficient Fine-tuning of Quantized LLMs” paper

We’ve covered the theory behind model quantization; now, let’s apply it. Let’s look at an example of how we can directly apply these quantization techniques to optimize our model, ensuring scalability and cost-efficiency without sacrificing performance using Hugging Face’s bitsandbytes.

Practical Implementation

💡 You can access the complete colab Notebook for this article here.

In practical settings, you don’t have to quantize models from scratch. Fortunately, several libraries and tools make quantizing models much easier. In our example, we use the bitsandbytes quantization library, which is based on LLM.int8(). This library reduces the precision of model weights by converting them from formats like FP16 or FP32 into lower-bit representations, typically 8-bit or 4-bit, thereby saving memory and speeding up computations without a substantial loss in performance. In this tutorial, we load a full-precision (FP16) Llama 3 model, convert it into a 4-bit quantized version, and compare their memory usage, generation quality, and speed.

Step 1: Setting Up the Environment

First, we set up the necessary imports and configure the cache directory. We import PyTorch for the deep learning framework, transformers for the model libraries, and utilities for memory management and timing. You need to also set up an access token to access private models on HuggingFace. You can create an access token here. You also need to ensure you have access to the meta-llama/Meta-Llama-3-8B model on HuggingFace since it’s a gated model.

import torch
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer, 
    BitsAndBytesConfig
)
import time
import gc
import os

# Set up your Hugging Face access token if you're using a private model
os.environ["HUGGINGFACE_HUB_TOKEN"] = "your_huggingface_token_here"

# Configure a custom cache directory to avoid 
# re-downloading large files.
CACHE_DIR = "/cache_dir/path"
os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
os.environ["HF_HOME"] = CACHE_DIR

We specify a cache directory to store downloaded models, which helps avoid repeatedly downloading large files when working with these models.

Step 2: Loading the Full Precision (FP16) Model

In this step, we define a function to load the full-precision model in FP16 format using the AutoModelForCausalLM class from the transformers library with the torch_dtype set to float16. This class automatically selects the appropriate model architecture based on the model identifier provided. We also load the corresponding tokenizer with AutoTokenizer.

def load_fp16_model():
    print("\n=== Loading Full Precision (FP16) Model ===")
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Meta-Llama-3-8B",
        torch_dtype=torch.float16,
        device_map={"": 0}
    )

    tokenizer = AutoTokenizer.from_pretrained(
        "meta-llama/Meta-Llama-3-8B"
    )

    # Calculate memory usage
    memory_gb = torch.cuda.max_memory_allocated() / (1024**3)
    print(f"Memory usage (FP16): {memory_gb:.2f} GB")

    return model, tokenizer, memory_gb

We specify that the model should be placed on the first GPU using the device_map parameter. We also load the corresponding tokenizer, which handles text-to-token conversions. We also calculate the GPU memory usage to understand how much VRAM the model consumes.

Step 3: Loading the 4-bit Quantized Model

Next, we define a function to load the 4-bit quantized version of the same model. Before loading, we clear the GPU memory to ensure accurate measurement. We configure quantization using BitsAndBytesConfig.

def load_4bit_model():
    print("\n=== Loading 4-bit Quantized Model ===")
    # Clear memory first
    gc.collect()
    torch.cuda.empty_cache()

    # Configure 4-bit quantization
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True
    )

    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Meta-Llama-3-8B",
        quantization_config=quantization_config,
        device_map={"": 0}
    )

    tokenizer = AutoTokenizer.from_pretrained(
        "meta-llama/Meta-Llama-3-8B"
    )

    # Calculate memory usage
    memory_gb = torch.cuda.max_memory_allocated() / (1024**3)
    print(f"Memory usage (4-bit): {memory_gb:.2f} GB")

    return model, tokenizer, memory_gb

We configure the 4-bit quantization using the BitsAndBytesConfig class, specifying that we want to load the model in 4-bit precision with the “nf4” quantization type (normalized float 4-bit). We enable double quantization for additional memory savings. After loading, we calculate the GPU memory usage to compare with the full-precision model.

Step 4: Creating a Text Generation Function

We then create a function to handle text generation for both models. The function tokenizes the input prompt, transfers it to the model’s device, and measures the generation time.

def run_generation(model, tokenizer, prompt):
    input_ids = tokenizer(
        prompt, return_tensors="pt"
    ).input_ids.to(model.device)

    # Time generation
    start_time = time.time()
    output = model.generate(
        input_ids,
        max_new_tokens=50,
        do_sample=True,
        temperature=0.7
    )
    generation_time = time.time() - start_time

    # Decode output tokens to text
    output_text = tokenizer.decode(output[0], skip_special_tokens=True)

    return output_text, generation_time

This function generates text with up to 50 new tokens using temperature sampling (0.7) for a balance of creativity and coherence. After generation, we decode the output tokens back to the text and return both the generated text and the time taken.

Step 5: Comparing the Models Qualitatively

Here, we define a function to compare the output quality of both the FP16 and 4-bit models using example prompts.

def qualitative_comparison(
    model_fp16, 
    model_4bit, 
    tokenizer, 
    example_prompts
):
    print("\n=== Qualitative Comparison: FP16 vs 4-bit ===")
    print("=" * 50)

    # Create file for saving full outputs
    with open("comparison_results.txt", "w") as f:
        f.write("=== Qualitative Comparison: FP16 vs 4-bit ===\n")

        for i, prompt in enumerate(example_prompts):
            print(f"\nExample {i+1}: \"{prompt}\"")
            print("-" * 50)

            # Generate with FP16
            fp16_output, fp16_time = run_generation(
                model_fp16, tokenizer, prompt
            )

            # Generate with 4-bit
            q4_output, q4_time = run_generation(
                model_4bit, tokenizer, prompt
            )

            # Print truncated results to console
            print(f"FP16 ({fp16_time:.2f}s): {fp16_output[:150]}...")
            print(f"4-bit ({q4_time:.2f}s): {q4_output[:150]}...")

            # Write full results to file
            f.write(f"\nExample {i+1}: \"{prompt}\"\n")
            f.write("-" * 50 + "\n")
            f.write(f"FP16 ({fp16_time:.2f}s):\n{fp16_output}\n\n")
            f.write(f"4-bit ({q4_time:.2f}s):\n{q4_output}\n\n")

For each prompt, we generate responses using both the FP16 and 4-bit models, recording the time taken for each generation. We print truncated outputs to the console for quick review and save the full outputs to a file for more detailed analysis later. This allows us to assess the models’ quality and speed differences.

Step 6: Running the Complete Comparison

Finally, we implement the main function that runs the entire comparison process. This includes memory measurement, model loading, and qualitative output comparisons.

if __name__ == "__main__":
    example_prompts = [
        "A robot discovers what it means to be human when",
        "Explain quantum computing to a 5-year old child:",
        "Write a short poem about artificial intelligence:",
        "The main difference between supervised and unsupervised learning is",
        "Summarize the plot of Romeo and Juliet in three sentences:"
    ]

    # Completely reset before measuring each model
    torch.cuda.empty_cache()
    gc.collect()
    torch.cuda.reset_peak_memory_stats()

    # Load FP16 model and measure
    fp16_model, tokenizer, fp16_memory = load_fp16_model()

    # Delete FP16 model before loading 4-bit model
    del fp16_model
    torch.cuda.empty_cache()
    gc.collect()
    torch.cuda.reset_peak_memory_stats()

    # Now load 4-bit model and measure
    q4_model, tokenizer_4bit, q4_memory = load_4bit_model()

    # Print memory usage stats
    memory_reduction = (fp16_memory - q4_memory) / fp16_memory * 100
    print(f"\nMemory usage comparison:")
    print(f"FP16: {fp16_memory:.2f} GB")
    print(f"4-bit: {q4_memory:.2f} GB")
    print(f"Reduction: {memory_reduction:.2f}%")

    # reload the FP16 model for the comparison
    torch.cuda.empty_cache()
    gc.collect()
    fp16_model, _, _ = load_fp16_model()

    # Run qualitative comparison
    qualitative_comparison(
        fp16_model, q4_model, tokenizer, example_prompts
    )

    # Clean up
    del fp16_model, q4_model
    torch.cuda.empty_cache()

We define test prompts covering various types of generation tasks, from creative writing to factual knowledge. We carefully manage GPU memory between model loads to ensure accurate measurements. We load the FP16 model first, measure its memory usage, then load the 4-bit model and measure its usage. We calculate and report the memory reduction achieved through quantization. We then reload the FP16 model (since we had to free its memory earlier) and run the qualitative comparison between both models, concluding by cleaning up resources.

Results

After running the complete comparison, you might see outputs like:

=== Loading 4-bit Quantized Model ===
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████| 4/4 [00:07<00:00,  1.89s/it]
Memory usage (4-bit): 5.42 GB

Memory usage comparison:
FP16: 14.96 GB
4-bit: 5.42 GB
Reduction: 63.73%

=== Loading Full Precision (FP16) Model  ===
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████| 4/4 [00:07<00:00,  1.77s/it]
Memory usage (FP16): 20.27 GB

=== Qualitative Comparison: FP16 vs 4-bit ===
==================================================

Example 1: "A robot discovers what it means to be human when"
--------------------------------------------------
FP16 (1.84s): A robot discovers what it means to be human when it falls in love with a human woman in this futuristic science fiction tale. The year is 2036 and the...
4-bit (2.06s): A robot discovers what it means to be human when it's forced to interact with real people. When it becomes clear that a nuclear strike is imminent, th...

Example 2: "Explain quantum computing to a 5-year old child:"
--------------------------------------------------
FP16 (1.40s): Explain quantum computing to a 5-year old child: A Q&A with IBM’s Dr. Talia Gershon
Dr. Talia Gershon is a quantum computing scientist at IBM. She has...
4-bit (2.03s): Explain quantum computing to a 5-year old child: a new way to think about quantum computing
Quantum computing is a very difficult concept to explain. ...

Example 3: "Write a short poem about artificial intelligence:"
--------------------------------------------------
FP16 (1.39s): Write a short poem about artificial intelligence: a poem about artificial intelligence.
A poem about artificial intelligence. A short poem about artif...
4-bit (2.02s): Write a short poem about artificial intelligence: The AI poem generator
AI is an amazing technology that can help us solve many problems. However, it’...

Example 4: "The main difference between supervised and unsupervised learning is"
--------------------------------------------------
FP16 (1.40s): The main difference between supervised and unsupervised learning is that supervised learning uses labeled data, while unsupervised learning uses unlab...
4-bit (2.03s): The main difference between supervised and unsupervised learning is that supervised learning is the learning process in which we have labeled data, wh...

Example 5: "Summarize the plot of Romeo and Juliet in three sentences:"
--------------------------------------------------
FP16 (1.39s): Summarize the plot of Romeo and Juliet in three sentences: Act 1
In Act I, Romeo and Juliet meet at a ball, fall in love, and decide to get married. T...
4-bit (2.04s): Summarize the plot of Romeo and Juliet in three sentences: What is the basic idea of Romeo and Juliet?
The basic idea of Romeo and Juliet is that two ...

Results Analysis

Looking at the outputs of FP16 and 4-bit quantized versions of the Llama 3 model, we can see a significant trade-off in memory usage, performance speed, and output quality. Quantization to 4-bit results in substantial memory efficiency, reducing the memory footprint from 14.96 GB in the FP16 model to 5.42 GB, representing approximately a 64% decrease. This considerable reduction is particularly advantageous for deployments with constrained memory resources.

However, we can see that quantization introduces a consistent slowdown in inference speed, with the 4-bit model experiencing around 30-45% longer response times compared to the FP16 model. Specifically, the 4-bit model typically generates outputs within 2.02-2.47 seconds, whereas the FP16 model completes similar tasks within 1.40-1.84 seconds. This slowdown is partly because we ran this setup on an NVIDIA GPU, which is highly optimized for FP16 computations through Tensor Cores. In contrast, native support for 4-bit operations on these GPUs is limited, resulting in additional overhead for dequantizing and scaling values during inference.

In terms of output quality, both models produced coherent and accurate responses across a variety of prompt types. While subtle differences are seen, such as the FP16 model occasionally provided more detailed answers or specific references, whereas the 4-bit model offered more generalized explanations. There was no significant degradation in quality. Structured creative tasks, like poetry and concise summarization, were challenging for both models, indicating quantization did not disproportionately impact performance in these areas.

Overall, the quantization process achieved notable memory savings without substantial compromise in the coherence or quality of the generated outputs, making it highly suitable for memory-sensitive applications where slight increases in latency are acceptable.

With these results in mind, let’s now focus on the key performance trade-offs and challenges that arise in real-world deployments, where balancing efficiency, accuracy, and response times becomes essential.

Performance Trade-offs and Challenges

Quantization generally enhances speed and memory efficiency, yet it introduces several challenges that require careful consideration:

Accuracy Degradation: Lowering precision generally impacts metrics such as accuracy. For instance, FP16 precision typically results in negligible loss, while well-implemented int8 quantization often leads to less than a 1% drop. In contrast, 4-bit quantization may noticeably degrade performance unless advanced methods like GPTQ or AWQ are applied. A practical strategy is to begin with 8-bit quantization and only experiment with 4-bit if further compression is needed while monitoring accuracy. If losses are excessive, consider quantization-aware training (QAT) or alternative schemes (e.g., per-channel scaling or SmoothQuant).
Selecting the Appropriate Method: Different scenarios call for different approaches. For quick CPU improvements, post-training quantization (PTQ) to int8 offers an easy 2–4× speed boost. On GPUs with limited memory, using 8-bit quantization (via options like load_in_8bit) is a solid choice, while 4-bit methods such as GPTQ or AWQ can compress models further without extra training. If fine-tuning is possible, QAT may slightly improve accuracy over PTQ. Also, the impact of quantization can vary by task—for instance, slight decreases in perplexity might more noticeably affect generative text quality than classification accuracy.
Benchmarking and Performance Gains: It’s essential to assess both accuracy and latency. Quantization yields benefits only when the hardware and runtime are optimized for lower precision. With specialized runtimes (e.g., IPEX, TensorRT), researchers have reported up to 3.5× faster inference on A100 GPUs using 3-bit GPTQ models and even 4.5× on older GPUs. On CPUs, int8 quantization can offer 4–8× speed improvements over FP32.
Integration with Other Compression Techniques: Quantization can be effectively combined with methods like pruning, distillation, and efficient architectures. For example, distilling a large model into a medium one and then applying int8 quantization produces a lightweight yet high-performing model. Integrated frameworks like the Intel Neural Compressor combine both pruning and quantization, although each added compression step requires careful evaluation to balance performance and accuracy.

In summary, while quantization dramatically improves performance, careful calibration, minimal fine-tuning, and advanced algorithms are crucial to maintaining accuracy. Always evaluate the quantized model using key metrics to ensure it meets your requirements.

Conclusion

Model quantization is a transformative technique for optimizing LLMs. By reducing numerical precision—from FP32 to formats such as FP16, INT8, or even 4-bit representations—quantization substantially lowers memory usage and computational demands. This process, whether through post-training quantization or quantization-aware training, enables faster inference and cost-effective deployment on resource-constrained devices while preserving acceptable accuracy. However, striking the right balance between efficiency and performance remains essential, as aggressive quantization can introduce trade-offs such as minor accuracy losses or slower generation speeds in some cases.

If you're working with large-scale LLMs, understanding techniques like quantization is just one piece of the puzzle. Optimizing model performance, fine-tuning effectively, and managing deployment trade-offs are all critical to building efficient AI systems. If you want to develop scalable, high-performance LLM products without wasting time on trial and error, our Beginner to Advanced LLM Developer course provides the in-depth guidance you need. Learn model optimization, fine-tuning strategies, and practical implementation techniques—all designed to help you build smarter, more efficient AI solutions.

How To Construct Self-Attention Mechanisms For Arbitrary Long Sequences

Damien Benveniste — Fri, 21 Mar 2025 15:01:59 GMT

With Gemini models having a 2M tokens context size and Claude having a 200K tokens context size while having a time-to-first-token this fast, it is impossible not to have a modification of the attention mechanism that explicitly handles extremely long sequences. Considering the size of those models, 2M tokens would lead to Petabytes of GPU memory to generate each token! So, we need strategies to handle those sequence sizes efficiently! Let’s dive in:

Transformer-XL
Memorizing Transformers
Infini-Attention

In the previous newsletters, we examined methods to reduce attention's computational complexity. In this newsletter, we are going to focus on designing attention mechanisms specifically optimized for processing extremely long contexts. The fundamental difference lies in the objective. Low-complexity attention methods primarily aim to approximate standard attention more efficiently, whereas long-sequence attention mechanisms fundamentally rethink how information flows across distant positions. Rather than merely making attention more computationally feasible, we are going to look at strategies to make distant contextual information meaningfully accessible and useful to the model.

Transformer-XL

Transformer-XL was proposed in 2019 as a way to process sequences of virtually unlimited length while maintaining coherent information flow across the entire document. The main limitations to handle sequences of any length are:

The typical time complexity O(N²) of the attention layers. However, we have already seen strategies to reduce that complexity to O(N). For an autoregressive process, we are bounded from below by, at least, a linear decoding process in the sequence length. Therefore, we can never do better than this theoretical constraint.
The absolute positional encoding proposed in the original "Attention is all you need" paper is the main blocker for encoding arbitrary sequence lengths. We would need a way to encode any possible positions, which is hard in practice.
Another important blocker is the memory constraint. Longer sequences take more space in memory, and we would reach an upper bound in length when the GPU memory becomes saturated. We saw when discussing the FlashAttention that a high-end NVIDIA A100-8GB could realistically handle a maximum sequence length of 5,932 for a GPT-3 model.

With Transformer-XL, we are going to design an attention mechanism that processes sequences with linear time complexity, constant memory complexity, and a novel relative positional encoding that captures the relative distance between tokens instead of their absolute positions. We are going to delay diving into the relative positional encoding until a future newsletter and focus here on how to process arbitrary sequence lengths in linear time with bounded memory constraints.

To illustrate how it works, let's consider the following toy example of input sequence:

"Teaching computers to see the world makes every colorful dataset an adventure"

The strategy with Transformer-XL is to handle the incoming tokens by segments. We break down the incoming sequence into segments, typically of size ~128-512 tokens during training and up to 1600 tokens during evaluation. For our toy example, let's assume that our segments are four tokens long:

Segment 1: ['Teaching', 'computers', 'to', 'see']
Segment 2: ['the', 'world', 'makes', 'every']
Segment 3: ['colorful', 'dataset', 'an', 'adventure']

Formally, we divide the incoming sequence into T = N / n segments, where n is the number of tokens per segment. We are going to create a segment-level recurrence to generate the output vectors from the model:

Generate all the hidden states [H¹₁, H²₁, …,H^L₁] related to the first segment 𝝉 = 1. l ∈ {1, …, L} is the index of the layer l, and L is the total number of layers in the model. The full attention is computed within segments, so its time and memory complexity is O(n²). We are going to cache the intermediate representations of the tokens [H¹₁, H²₁, …,H^L₁] for the next iteration.
In the next iteration, we are going to consider the second segment and its interaction with the first segment. At each layer, we retrieve the hidden states of segment 𝝉 = 1 and append them to the hidden state of segment 2
and we compute the next hidden states by passing them through the layer:
Here, the attention matrix is computed across segments 1 and 2 and is, therefore, of size 2n x 2n, which still follows a quadratic complexity ~O(n²).
In general, for any segment 𝝉, we retrieve the computed hidden states [H¹_𝝉-1, H²_𝝉-1, …,H^L_𝝉-1] for the previous segment 𝝉 - 1, and append them to the hidden states of the current segment H^l_𝝉 for layer l:
and compute the next hidden states:
At every point during this recurring process, the time and space complexity is at most ~O(n²), and we iterate this process until we reach the last segment in the sequence.

The recurrence mechanism effectively creates a form of "memory" that allows the model to maintain coherent understanding across very long texts while keeping computational requirements manageable. Transformer-XL's ability to maintain coherence across long contexts is directly related to its depth. H^l+1_𝝉 depends on H^l_𝝉-1 and H^l_𝝉, which means that the generation of the n hidden states in H^l+1_𝝉 depends on the 2 x n hidden states in [H^l_𝝉-1; H^l_𝝉]. The hidden states H^l_𝝉-1 also depend on H^l-1_𝝉-2 and H^l-1_𝝉-1. Therefore, H^l+1_𝝉 depends on the 3 x n hidden states in H^l_𝝉, H^l-1_𝝉-1, H^l-1_𝝉-2. If we go back k layers in the model, then the hidden states in H^l+1_𝝉 depend on the (k+1) x n hidden states in [H^l-k+1_𝝉, H^l-k+1_𝝉-1, …, H^l-k+1_𝝉-k]. By unrolling the recurrence of the hidden states' dependency, we establish the functional relationship:

This creates a recurrence relation where information from segment 𝝉 −k can reach the final layer L in segment 𝝉 only if k+1 ≤ L. If we want to cover all the N tokens in the sequence, with k = N/n - 1, this means we need:

The theoretical dependency length is, therefore, O(n x L). If the network is not deep enough, its ability to propagate information across segments would be limited. This is because information flows through layers within a segment before being passed to the next segment. A shallow Transformer-XL would still have the unbounded context mechanism but might not effectively utilize information from distant parts of the sequence.

So far, we have only considered the forward pass, but it becomes messy when we consider the backward pass! When we compute the gradient of the loss function L, we need to compute its relation to the hidden states H^l+1_𝝉:

The layer-wise summation

must consider all layers from l-k+1 up to l+1 and the segment-wise summation

must consider all segments from 𝝉 −k up to 𝝉. This creates (k+1)² gradient paths for a single hidden state segment. Now, for a sequence of length N divided into N/n segments, if we extend this to consider all possible hidden states that influence the loss, we get approximately:

gradient paths to compute. This cubic growth makes training infeasible for long sequences. We are going to modify the segment-level dependency by introducing the Stop-Gradient SG(.) operator (detach in PyTorch):

This operation prevents gradients from flowing backward through the cached states during backpropagation. This means:

By applying SG(.), the previous segment hidden states H^l_𝝉-1 is treated as constant during the backward pass. When we apply SG(.) to prevent gradient flow across segment boundaries, we introduce a profound asymmetry:

Information Asymmetry: During the forward pass, the model can access and use information from previous segments, but during backpropagation, it cannot receive gradient signals from future segments.
Truncated Credit Assignment: The model cannot directly attribute credit or blame to decisions made in previous segments, even though those decisions influence future outcomes.

This creates a unique learning paradigm where the model must learn to encode useful information in its hidden states without direct optimization signals for long-range dependencies. Interestingly, despite not being directly optimized for very long dependencies, Transformer-XL still develops impressive long-range capabilities because the recurrence mechanism creates paths for information to propagate forward.

For any two layers a and b, we have:

Therefore, the loss backpropagation simplifies to

which reduces the computational complexity from the cubic behavior O((N/n)³) to the linear one O(N/n) .

Each segment has length n, and when processing a single segment, attention is computed over the current segment plus the cached previous segment. The time complexity of the computation per segment is, therefore, ~O(4n²) = O(n²). For a sequence of length N, we process approximately T = N/n segments, and each segment requires O(n²). Therefore, the complexity of the total operations is:

Which is linear in the sequence size N. When it comes to the way the attention is computed, it is a very similar pattern to the sparse attention with a sliding window. Beyond the current and past segments, the hidden states are blind to the other tokens.

However, with sparse attention, the memory requirements grow linearly with the sequence size, but with Transformer-XL, the space complexity is bounded by O(n²) = O(1) as n is a fixed number. We still need to choose n large enough to make efficient use of the high GPU parallelism.

Memorizing Transformers

In Transformer-XL, the long-range coherence is captured by the successive layers in the model, but the direct interaction between tokens is lost beyond a two-segment window. Google introduced the Memorizing Transformers in 2022 that extended the long-range coherence by caching previous key-value pairs for selective retrieval. The attention computation is broken down into two parts:

The local attention: As before, we partition the input sequence into segments (usually 512 tokens) and compute the token-token interactions within each segment:
where Q_local, K_local, and V_local are the local queries, keys, and values within the segment.

How To Improve Decoding Latency With Faster Self-Attention Mechanisms

Damien Benveniste — Wed, 12 Mar 2025 15:02:19 GMT

In LLMs, handling large sequences is not enough, we need to make sure the decoding process is fast. Here we explore 3 typical approaches used to speed up the decoding process:

Multi‑Query Attention
Grouped‑Query Attention
DeepSeek Multi-head Latent Attention

When generating text one token at a time, transformers face a significant performance challenge. In the original transformer architecture with Multi-Head Attention (MHA), each time we generate a new token, we need to:

Load the entire history of key (K) and value (V) matrices for each attention head
Process them against the new query
Generate the next token
Repeat

This creates a memory bandwidth bottleneck. For long sequences, we're constantly reloading massive tensors from memory, which becomes the limiting factor in generation speed. Specifically, with n_head attention heads and a sequence of length N, the original approach required loading tensors of size approximately n_head x d_model x N, which could be gigabytes of data for long contexts. Since those K and V tensors are a bottleneck for the decoding process, we are going to look at strategies to minimize their memory requirements.

Strategies like multi‑query attention, grouped-query attention, and multi-latent attention are highly coupled with the KV-caching technique. With KV-caching, instead of recomputing the same K and V matrices over and over, we cache them in memory and update them at each iteration. Loading the KV-cache from memory becomes the dominant bottleneck instead of the compute itself.

Multi‑Query Attention

In 2019, Noam Shazeer published a variant of the multi-head attention: the multi-query attention (MQA). He realized that the model could maintain most of its capability while sharing a single set of keys and values across all heads. This insight might seem counterintuitive since the whole point of multi-head attention was to have different "perspectives" on the same information. What Shazeer discovered was that the diversity in the query projections still allows different heads to extract different information, the shared key-value store acts as a common knowledge repository, and each head can still attend to different parts of this shared repository.

Formally, this means that the projection matrix W^Q is still a d_model x d_model matrix, but the projections W^K and W^V are d_model x d_head matrices. With an incoming hidden state H of dimension d_model x N, we have at training time the initial projections:

We then reshape the matrices into tensors to highlight the number of heads:

At inference time, we only consider the last query q’_N of size d_head x n_head in the input sequence since we only need to predict the last token. The alignment scores are computed by broadcasting the matrix K’ to all the heads:

We perform the softmax transformation:

The value is again broadcasted to all heads to compute the context vector corresponding to the prediction:

The context vector is reshaped in the original dimensions:

Finally, the context vector is projected one last time to mix the information from the different heads:

Let's consider the memory access complexity (or memory bandwidth complexity). It measures the total amount of data that must be transferred between memory and compute units during the entire sequence of operations. This measures bandwidth requirements. For MHA, at each decoding step, we need to load W^Q, W^K, W^V, and W^O. They are all d_model x d_model matrices, so over N decoding steps, the memory access complexity is ~O(Nd_model²). As we generate each token, we must reload the entire history. For the i-th token, we load keys and values of size i x d_model. Summing over all N steps:

We also need to load the N input hidden states of size d_model. So, the overall memory access complexity for MHA is

For long sequences, d_modelN² dominates, which is the problematic bottleneck.

For MQA, loading W^Q, W^K, W^V, and W^O is the same asymptotic behavior ~O(Nd_model²) and the input hidden states as well ~O(Nd_model). However, the K and V matrices have a size i x d_head at the i-th decoding step, which leads to:

So, the overall memory access complexity is

As the d_headN² term dominates for long sequences, it has a substantial impact because memory access increase leads to higher latency. For the specific experiments run in Shazeer's paper, he found that MQA achieves ~12x faster decoder inference with minimal performance loss. For inference with a sequence length of 128 tokens, MHA's decoding time per token was 46 μs whereas MQA's was 3.8 μs.

Grouped‑Query Attention

With multi-query attention, we gain in decoding speed, but we lose in performance compared to multi-head attention. The grouped-query attention provides a middle ground between MQA and MHA to keep performance high while improving decoding speed performance. GQA divides query heads into G groups, where each group shares a single key head and value head. This creates a configurable spectrum:

When G = 1, it is equivalent to MQA (single key-value head for all queries)
When G = n_head, it is equivalent to standard MHA (separate key-value for each query)
1 < G < n_head, we have the GQA sweet spot that balances efficiency and quality

It is important to understand that latency is not proportional to memory access. Memory-level parallelism refers to a computer system's ability to process multiple memory operations simultaneously rather than sequentially. There is a regime where the memory-level parallelism of the hardware is underutilized when we are trying to load too few groups, and there is a regime where we are saturating it when we are loading too many groups. Because of the parallelism, it is quite likely that having G = 1 (MQA) induces as much latency as G = 4 despite the increased predictive performance. This relationship between groups and latency is highly hardware-specific, which is why empirical testing is necessary to find the optimal configuration for any given system. The TPUs used in the original paper showed this saturation around G = 8, but different architectures might have different saturation points.

From https://arxiv.org/pdf/2305.13245

GQA achieves quality close to MHA with speed comparable to MQA. GQA is more stable during training than pure MQA, which can sometimes exhibit training instability. For example, on the T5-XXL model, the GQA-8 version presented in the original paper achieved an average performance of 47.1 across key benchmarks compared to 47.2 for MHA and 46.6 for MQA, while maintaining inference speeds much closer to MQA.

From MHA to GQA

One of the main appeals of GQA is the ability to convert a vanilla MHA into GQA with minimal effort. The ability to convert existing models rather than train new ones has accelerated the adoption of GQA across the industry. Large models like PaLM 2 and LLaMA 2 have incorporated GQA, partly because this conversion process made it practical to do so. Converting from MHA to GQA is a two steps process:

Checkpoint Conversion: We first convert the model's weights by mean-pooling the key and value projection matrices within each group.
Additional Pre-training: After conversion, we continue pre-training the model for a small fraction (about 5%) of the original training steps, using the same pre-training dataset and objectives.

This approach requires only about 5% of the original training compute to achieve comparable performance to the original MHA model. For large models that cost millions to train, this represents enormous savings.

DeepSeek’s Multi-head Latent Attention

Multi-Head Latent Attention (MLA) was introduced in DeepSeek-V2 as a way to optimize for training and inference speed while preserving the predictive performance. Even if GQA keeps performance high, it still has to balance efficiency with performance. MLA provides a solution that improves on both fronts at the same time.

How To Reduce The Memory Usage Of The Self-Attention

Damien Benveniste — Thu, 06 Mar 2025 16:03:31 GMT

With a bit of magic, we take a very inefficient computation like the Self-Attention and make it super memory-optimized for the specific hardware we use for training and inference. And we all need a bit of magic!

Self-attention Does Not Need O(N²) Memory
The GPU Architecture
The FlashAttention-1
The FlashAttention-2
The FlashAttention-3

So far, we have mainly explored how to reduce the complexity of the attention mechanisms by approximating the vanilla attention. The vanilla attention has a strict O(N²) time complexity, but the space complexity doesn't need to be O(N²)! Computing Q^TK requires ~O(N²) operations, but the full N x N alignment scores and attention matrices do not need to be fully materialized all at once in memory. As models and sequence lengths scale, it becomes essential to minimize the memory requirements at training and inference time to better utilize the underlying hardware.

Self-attention Does Not Need O(N²) Memory

Let's consider again the computation of the context vectors:

Here is how we arrive at the typical O(N²) space complexity:

The typical assumption is that we first compute the dot product between the query q_i and all the keys [k₁, …,k_N]:
Where e_i is the alignment score vector of size N for the query q_i, which leads to the N x N matrix for the N queries.
We then perform the softmax transformation:
Here, a_i is the attention vector of size N for the query q_i. Again, for N queries, it leads to the typical N x N attention matrix.
And finally, we project a_i onto the different values [v₁, …,v_N], which leads to c_i:

Therefore, naively computing the alignment scores and the attention matrices first forces the materialization of those matrices in memory, which leads to the O(N²) space complexity.

However, we do not need to order the computations in this manner! In 2021, Rabe and Staats realized that by reordering the operations, we can greatly reduce the requirements on the memory. The idea is to consider the unnormalized context vector c̃_i and the normalization constant ∑exp(q_i^Tk_j) of the softmax transformation separately:

Because c̃_i and s_i are just sums, we can easily loop through the key-value pairs to compute the context vector:

At any point during the for-loop, we only need to store the intermediary values of c̃_i and s_i. c̃_i is a vector of size d_model (ignoring heads for simplicity), and s_i is a scalar. Therefore, for one query, we need constant space complexity O(1) to compute one context vector. Even iterating through all the queries, we never need to capture more than the intermediary values of c̃_i and s_i, so we can compute the full attention mechanism in O(1) space complexity:

In reality, this is not a practical solution because sequential operations are not adapted to the parallelization capability of the CPU, GPU, or TPU hardware that is commonly used for neural network computations. In practice, the queries, keys, and values are partitioned into chunks to allow for a high degree of parallelization while keeping the memory requirement low. Let's assume that we partition the queries into n_q chunks and the keys and values into n_k chunks:

where each Q_i is a N / n_q x d_model matrix and K_i, V_i are N / n_k x d_model matrices. Let's call N_q= N / n_q, the number of queries per chunk, and N_k= N / n_k, the number of key-value pairs per chunk. We can now iterate through the chunks exactly in the same way:

As before, we need to store intermediary values of C̃_i and S_i. In this context, Q_iK^T_j is a matrix of size N_q x N_k, and so is A_ij. C̃_i is a matrix of size N_q x d_model and S_i is a vector of size N_q. Therefore the space complexity is O(N_q x N_k + N_q x d_model). To balance the number of chunks and the number of key-value pairs per chunk, they chose N_k= n_k= √N and fixed N_q= 1024. This results in a space complexity:

This approach allows for efficient tensor operations within each chunk while dramatically reducing the peak memory requirements. Note that no approximation has been made, and it is mathematically equivalent to the vanilla attention mechanism. However, this approach is slower (8-13% slower during the forward pass and 30-35% slower during the backward pass) due to the sequential computations, but it enables the processing of much longer sequences that would otherwise be impossible due to memory constraints. Along with the FlashAttention, it is one of the memory optimization strategies used in the xFormers package developed by Meta and used in the development of the Llama models.

Stabilizing The Computations

Until now, we have been ignoring the numerical stability of the softmax computation, but most implementations (PyTorch, TensorFlow, ...) are using a couple of tricks to ensure its stability. Let's remind ourselves of the softmax function:

Computing e^xi can be tricky because if x_i ≥ 89, then e^xi ≅ 4.4e38, which exceeds the floating-point limit of 3.4e38 for a 32-bit float number, potentially leading to float overflow errors. To prevent this from happening, we typically modify the exponent by finding the maximum x_i value:

How To Linearize The Attention Mechanism!

Damien Benveniste — Wed, 26 Feb 2025 16:01:43 GMT

Today, we talk about how to engineer attention mechanisms in O(n) complexity instead of O(n²). This newsletter tends to be a bit more math-flavored than my usual content, but it is liberating to be able to use math for the greater good!

Low-Rank Projection of Attention Matrices: Linformer
Recurrent Attention Equivalence: The Linear Transformer
Kernel Approximation: Performer

Self-attention’s quadratic complexity in sequence length has long been a central bottleneck for large-scale Transformer models. Handling tens of thousands of tokens becomes computationally prohibitive and can quickly exhaust available memory. Linear attention mechanisms represent a paradigm shift in transformer architecture by mathematically re-engineering the attention operation to achieve O(n) complexity while maintaining global context awareness. Unlike sparse attention's pattern restrictions, which preserve quadratic complexity but limit interactions to predefined token subsets, linear attention fundamentally redefines how all tokens interact by reformulating the attention matrix computation rather than pruning token interactions. Where sparse attention sacrifices theoretical completeness for practical speed, linear attention preserves global relationships at the cost of approximating pairwise token influences. This enables native handling of extreme sequence lengths (1M+ tokens) while avoiding sparse attention's blind spots.

Low-Rank Projection of Attention Matrices: Linformer

With Sparse attention mechanisms, we understood that most of the token interaction information was contained in a small subset of token pairs. Linformer introduced the idea that the token-token interaction matrix could be compressed into a smaller representation without too much information loss. Instead of computing the full N x N interaction Q^TK / √d (ignoring heads for simplicity), we could first project K into a lower rank dimension k, and compute the lower rank N x k approximation:

where E is a N x k projection matrix that projects K from the original dimension d x N to d x k. This leads to N x k alignment score and attention matrices.

When we project with E, the approximation leads to the error:

If the elements of E follow a Gaussian distribution ~N(0, 1/k), the Johnson–Lindenstrauss lemma guarantees that:

\\epsilon\\right]\\leq e^{-\\gamma\\epsilon^2 k}.","id":"VRYEADHXVE"}" data-component-name="LatexBlockToDOM">

This means that the probability that we choose E such that the error is greater than 𝛜 is bounded by exp(-𝝲𝛜²k), where 𝝲 is just a scaling constant. If we choose k → ∞, then P[error > 𝛜] → 0 for any 𝛜. A good choice is k ~ log N / 𝛜², yielding:

\\epsilon\\right]\\leq N^{-\\gamma}.","id":"XMEKFCFXZO"}" data-component-name="LatexBlockToDOM">

This means that we can choose an arbitrarily small 𝛜 such that P[error > 𝛜] → 0 as the sequence length increases N → ∞. Understand this as a mere theoretical guide that tells us that choosing k ~ log N will guarantee smaller errors as N increases. In practice, k is chosen independently of N, leading to the O(N) linear complexity while accepting the cost of the approximation error. Additionally, E is chosen as a parameter layer for the model to learn. For example, they showed that choosing k = 64 with N = 512 leads to slightly worse performance than the full attention.

Since the attention matrix has dimension N x k, we also need to project the values:

where F is the N x k projection matrix for the tensor V. As for E, F is also learned during training.

Projecting the keys and values EK, FV leads to complexity O(Nk). Computing the alignment scores Q^TEK and the context vectors C = AFV are also following O(Nk). Since we fix k, the overall time and space complexity is O(N).

Recurrent Attention Equivalence: The Linear Transformer

So far, we have accepted the attention mechanism to be represented by the following computation:

However, this specific analytical choice is not the only one that could be chosen to fulfill the same functional role in capturing pairwise interactions between tokens. Let's review the roles of the different elements in this equation:

The dot-product Q^TK: Similarity computation. For each query vector, it tells you how "compatible" or similar it is to each key vector. This yields a matrix of unnormalized attention scores.
Normalizing by √d: Variance control. The primary purpose of scaling by √d is to control the scale of the attention logits before softmax, ensuring stable gradient flow and preventing the softmax from becoming too "confident" (peaked). Furthermore, extremely large logits can cause numerical instability (e.g., NaN in floating-point arithmetic), and scaling mitigates this.
Softmax operation: Normalization and nonlinearity. The softmax turns the unnormalized similarity scores into a probability distribution, amplifying the effect of the most relevant keys.
Multiplication by V: Weighted aggregation. Each output is a weighted sum of the values, where the weights come from the normalized similarity scores. This is how the model “mixes” information from across the input sequence.

Functionally, we need a similarity function sim that is non-linear and captures the pairwise token interaction:

where the denominator ensures that the similarity function is normalized to 1. If we choose sim(q_i, k_j) = exp(q_i^Tk_j / √d), we recover the softmax transformation. The Linear Transformer proposed a new attention mechanism with a different analytical form, but with similar functional roles. More specifically, they suggested a similarity function where we can factorize the contribution from the keys and the queries as a product:

In the context of kernel methods in machine learning, ɸ is called a "feature map". A feature map is a function that transforms an input vector into a new space, often a higher-dimensional one, so that a kernel function (which measures similarity) can be expressed as an inner product in that space. Essentially, ɸ extracts or "maps" the original features into a new representation where the desired similarity (that mimics the softmax behavior) is computed simply by taking a dot product. In the context of the Linear Transformer, they simply chose ɸ as follows:

0, \\\\\n\n\\exp x & \\text{otherwise}.\n\n\\end{cases}","id":"FAIVUTXUGR"}" data-component-name="LatexBlockToDOM">

This ensures that sim(q_i, k_j) is always positive and is computationally stable while being non-linear. The main appeal of this linearization of the similarity kernel is the associativity property of the matrix multiplication:

For one key and one query, ɸ(q_i)^Tɸ(k_j) takes d operations. Multiplying the resulting scalar alignment score to v_j takes another d operations. Therefore, for all the keys, computing ∑ɸ(q_i)^Tɸ(k_j)v_j takes 2Nd operations, and the time complexity is O(Nd) per query. Similarly, the denominator ∑ɸ(q_i)ɸ(k_j) follows a time complexity of O(Nd). Because we have N queries, the total cost is

If we consider the multiplications in a different order, ɸ(k_j)v_j^T is an outer product and results in d² operations. For N keys and values, we end up with Nd² operations for ∑ɸ(k_j)v_j^T. In the denominator, summing the different keys ∑ɸ(k_j) requires Nd operations. Let's call S = ∑ɸ(k_j)v_j^T and z = ∑ɸ(k_j). S is a matrix of size d x d and z is a vector of size d. Computing ɸ(q_i)^TS brings another d² operations, and computing ɸ(q_i)^Tz takes d operations. Therefore, the cost of ɸ(q_i)^TS / ɸ(q_i)^Tz per query is O(d² + d) = O(d²). For N queries, we obtain a total complexity of:

Understanding The Sparse Transformers!

Damien Benveniste — Wed, 19 Feb 2025 16:01:50 GMT

The First Sparse Attention: Sparse Transformers
Choosing Sparsity Efficiently: Reformer
Local vs Global Attention: Longformer and BigBird

The original Transformer architecture introduced in the "Attention is All You Need" paper opened a whole new avenue of research, and it became useful to address some of the bottlenecks that this architecture came with. Today, modern Transformers rarely use the original vanilla Transformer blueprint without modifications. Instead, they often combine multiple techniques:

Faster or more memory-friendly attentions (sparse, linear, or memory-efficient) for large contexts,
Improved positional schemes (RoPE, ALiBi, or relative embeddings),
Enhanced feed-forward layers (MoE, GLU variants), and
Better normalization/optimizer choices (RMSNorm, AdamW).

These enhancements address the core bottlenecks, quadratic complexity, high memory usage, limited context, and potentially weak or rigid feed-forward/positional representations, making Transformers more scalable, expressive, and practical for today’s large language modeling tasks.

Sparse attention addresses this bottleneck by restricting, or "sparsifying", which tokens can attend to which. Instead of forming attention connections from every token to every other token, sparse mechanisms allow each token to attend to a subset of the sequence according to a specific pattern. By reducing the total number of key/value pairs, sparse attention can often achieve O(N log N) or even O(N) complexity.

The First Sparse Attention: Sparse Transformers

One of the first attempts at sparse attention was proposed by OpenAI in 2019, and this was the strategy chosen to build GPT-3. The idea is to limit the number of keys each query can attend when it comes to computing the alignment scores. This will reduce the number of alignment scores and attention weights computed. Because of this, we need to subset the values as well to ensure we can compute the context vectors C = AV^T.

OpenAI suggested two different sparse patterns where, in the different heads, the queries attend the keys differently. The first pattern is the strided pattern. One head focuses on a local window of nearby tokens by having the i-th query only attend the keys in [i-w, i], where w is the window size. For example, if w = 64, it means we only select the keys [i - 64, i - 63, …, i - 1, i]. The other heads focus on more global token interactions by having i-th query attending every c key. c is the stride and can be different for each head. For example, if c = 8, then we would only select the keys [0, …, i - 24,i - 16, i - 8, i].

They observed that the strided pattern (attending every k-th position) did not work well for text data, which lacks a naturally periodic structure. As a result, they introduced a fixed attention pattern. In one attention head, the sequence is divided into blocks of fixed size (e.g., 128 tokens). Each token within a block attends only to other tokens in that same block, capturing local dependencies in a more straightforward manner.

However, using purely local attention inside blocks would prevent information from flowing across blocks in deeper layers. To address this, another head connects the "summary token" at the end of each block to the corresponding summary tokens of all previous blocks. Because that last token has attended to all tokens in its own block, its hidden state acts as a summary of that entire sub-sequence. By letting future blocks attend to these summary tokens, the model propagates information globally across blocks, layer by layer.

Thus, the in-block pattern computes context vectors (weighted averages of values) only from nearby tokens, while the across-block pattern computes context from previous summary tokens. In combination, these patterns ensure both local and long-range context can be aggregated throughout the Transformer stack, even without a dense (i.e., fully quadratic) attention mechanism.

Let's estimate the time complexity of those sparse attentions. In the strided case, we have a local window of size w and a stride c. Therefore, each query attends to roughly w + N / c keys (the local window plus the strided tokens). For N queries, the total cost is ~ N(w + N / c). By tuning w and c, one can achieve sub-quadratic complexity. For example:

In the fixed case, The sequence is split into blocks of length l. Within each block, each query attends at most l keys. Furthermore, each query attends c = N / l summary tokens. Therefore, a query sees O(l + c) keys, and the total cost for all queries is O(N(l + c)). Typically, we choose l such that it grows sub-linearly with N. For example:

Despite the improved time complexity, it is essential for those computations to remain performed as tensor operations to fully utilize the high parallelism of the GPU hardware. For example, let's assume that we want to compute the attentions for the local sliding window described in the strided case. The original keys tensor K is of dimension n_head X d_head X N. We can construct the windowed keys tensor K^w to compute the sliding window all at once by adding another dimension representing the window size w. Constructing this tensor is a O(N) operation. The resulting tensor is of size n_head X d_head X N X w, and each slice of size d_head X N X w contains the necessary keys to compute the windowed attentions for each head.

Let's now compute the product between the windowed keys K^w and the queries Q:

where h represents the head dimension (n_head), n the sequence dimension (N), d the hidden size per head dimension (d_head) and w the window dimension. The resulting alignment scores tensor E is of shape n_head X N X w. The time complexity of this operation is O(d_headNw) instead of the vanilla O(d_headN²).

Choosing Sparsity Efficiently: Reformer

Attention Is All You Need: The Original Transformer Architecture

Damien Benveniste — Wed, 12 Feb 2025 16:02:31 GMT

This newsletter is the latest chapter of the Big Book of Large Language Models. You can find the preview here, and the full chapter is available in this newsletter

The Self-Attention Mechanism
The Multi-head Attention Layer
The Positional Encoding
The Encoder
The Residual Connections
The Layer Normalization
The Position-wise Feed-Forward Network
The Decoder
The Cross-Attention
Masking The Self-Attention Layer
The Prediction Head
The Decoding Process
Training For Causal Language Modeling
Understanding the scale of the model
Estimating The Number Of Model Parameters
Estimating The Floating‐Point Operations
The Different Architecture Variations
The Encoder-Only Architecture
The Decoder-Only Architecture
The Encoder-Decoder Architecture

The "Attention Is All You Need" paper is one of the most influential works in modern AI. By replacing recurrence with self-attention mechanisms, the authors introduced the Transformer architecture, a design that enabled parallelized training, captured long-range dependencies in data, and scaled effortlessly to unprecedented model sizes. This innovation not only rendered RNNs obsolete but also laid the groundwork for BERT, GPT, and the modern LLM revolution, powering breakthroughs from conversational AI to protein folding. Beyond technical innovations, the paper catalyzed a paradigm shift toward general-purpose models with the rise of foundation models trained on massive datasets and reshaped industries from healthcare to creative arts. In essence, it transformed how humanity interacts with language, knowledge, and intelligence itself.

Architecture Overview

The original Transformer architecture is composed of the encoder that computes a rich representation of the input sequence, the decoder that generates the output sequence, and the prediction head that uses the decoder output to predict the tokens of the output sequence.

The architecture presented in the "Attention Is All You Need" paper builds directly from the RNN encoder-decoder architecture while discarding recurrence entirely and replacing Bahdanau/Luong's cross-attention with intra-sequence attention. There are four important components to the architecture:

The embeddings: Besides the token embeddings necessary to project the tokens into their vector representations, the Transformer introduced the need for positional encoding to ensure that the information related to the token positions is captured by the model.
The encoder: As for the RNN encoder-decoder, the encoder is in charge of encoding the input sequence into vector representations such that the decoder has enough information to decode the output sequence. It comprises a stack of identical encoder blocks, each with multi-head self-attention (capturing global dependencies) and a position-wise feed-forward network (applying non-linear transformations).
The decoder: Similar to the encoder but adds masked multi-head self-attention (preventing future token visibility) and encoder-decoder attention (aligning decoder inputs with encoder outputs, akin to Bahdanau/Luong but without RNNs). As before, the autoregressive generation proceeds token-by-token.
Prediction head: The prediction head is a classifier over the whole token vocabulary made out of a linear layer followed by Softmax, converting the decoder's final hidden states into token probabilities to predict the next word.

We will cover each component in detail in the remainder of this chapter. Self-attention and position embedding are central to the transformer architecture, and we need to discuss those technical innovations before we can understand the entire architecture.

The Self-Attention

The Self-Attention Mechanism

The Architecture

In the case of the Bahdanau/Luong attention, the goal was to capture the interactions between the tokens of the input sequence and the ones of the output sequence. In the Transformer, the self-attention captures the token interactions within the sequences. It is composed of three linear layers: W^K, W^Q, and W^V. The input vectors to the attention layer are the internal hidden states h_i resulting from the model inputs. There are as many hidden states as tokens in the input sequence, and h_i corresponds to the i^th token. W^K, W^Q and W^V project the incoming hidden states into the so-called keys k_i, queries q_i and values v_i:

W^K, W^Q, and W^V are used to project the hidden states into keys, queries, and values.

The keys and queries are used to compute the alignment scores.

The keys and queries are used to compute the alignment scores:

As in the case of the Bahdanau attention, e_ij is the alignment score between the i^th word and the j^th word in the input sequence. d_model is the common naming convention for the hidden size:

The scaling factor √d_model in the scaled dot-product is used to counteract the effect of the dot product's magnitude growing with the dimensionality d_model, which stabilizes gradients and ensures numerical stability during training. It is common to represent to represent those operations as matrix multiplications. With the matrix K =[k₁, …, k_N] and Q =[q₁, …, q_N], we have:

or:

with N being the number of tokens in the sequence.

The attention weights are the result of normalizing the alignment scores by using the softmax transformation.

As for the other attentions, the alignment scores are normalized to 1 through a Softmax transformation:

where a_ij is the attention weight between the tokens i and j, quantifying how strongly the model should attend to token j when processing token i. Because we have Σ a_ij = 1, a_ij can be interpreted as the probability that token j is relevant to token i.

Each context vector is the result of a weighted average of the value vectors by using the attention weights.

The attention weights are used to compute a weighted average of the values vectors:

In the jargon used in the previous chapter, c_i are the context vectors coming out of the attention layer, but we can think of them as another intermediary set of hidden states within the network. Using the more common matrix notation, we have:

where V =[v₁, …, v_N], C =[c₁, …, c_N] and A = Softmax(E) is the matrix of attention weights.

The entire attention layer process.

The whole set of computations happening in the attention layer can be summarized as the following equation:

The Keys, Queries, and Values Naming Convention

The names "queries," "keys," and "values" are inspired by information retrieval systems (such as databases or search engines). Each token generates a query, key, and value to "retrieve" relevant context from other tokens. The model learns to search for relationships between tokens dynamically.

The queries represent what the current token is "asking for." For example, the word "it" in "The cat sat because it was tired," the query seeks antecedents (e.g., "cat"). The keys represent what other tokens "offer" as context. In our example, the key for "cat" signals it is a candidate antecedent for "it." The values are the actual content to aggregate based on attention weights. The value for "cat" encodes its contextual meaning (e.g., entity type, role in the sentence, ...). For each query (current token), the model "retrieves" values (context) by comparing the query to all keys (other tokens). For example, let us consider the sentence:

"The bank is steep, so it's dangerous to stand near it."

Query ("it"): "What does 'it' refer to?"
Keys ("bank," "steep," "dangerous"): Highlight candidates for reference.
Values: Encode the meaning of each candidate.

The model computes high attention weights between the query ("it") and keys ("bank," "steep"), then aggregates their values to infer "it" refers to the riverbank.

The Multi-head Attention Layer

The Naive Description

We have talked about self-attention so far, but we use the so-called multi-head attention layer in the transformer architecture. The multi-head attention layer works as multiple parallel attention mechanisms. By having multiple attention layers in parallel, they will be able to learn various interaction patterns between the different tokens in the input sequence. Combining those will lead to more heterogeneous learning, and we will be able to learn richer information from the input sequence. Think about the multi-head attention layer as an ensemble of self-attentions, a bit like the random forest is an ensemble of decision tree models.

We call "heads" the parallel attention mechanisms. To ensure that the time complexity of the computations remains independent of the number of attention heads, we need to reduce the size of the internal vectors within the layers. The hidden size dimensionality per head is divided by the number of heads:

where n_head is the number of heads. This implies that the hidden size has to be chosen so that it is divisible by the number of heads.

Each attention generates vectors of size d_head = d_model/ n_head, depending of the number of heads.

Let us call H =[h₁, …, h_N] the incoming hidden states. Each head h generates resulting hidden states H’_h of size d_head= d_model/ n_head:

To combine those heads' hidden states, we concatenate them, and we pass them through a final linear layer W^O to mix the signals coming from the different heads:

The result of each head is concatenated, and the signals are further mixed by a final linear layer W^O.

To generate smaller hidden states, we need to reduce the dimensionality of the internal matrices. In each head, the projection matrices W^K, W^Q, and W^V take vectors of size d_model and generate vectors of size d_model/ n_head.

To generate smaller context vectors, we underlying projection matrices W^K, W^Q, and W^V need to be of size d_model X d_head for each head.

The Tensor Representation

In reality, the projection matrices are not spread across multiple heads, but different sections of the matrices handle the projections for the different heads.

Although the information we have described so far about the multi-head attention layer is accurate, there is a critical subtlety to understand when it comes to its implementation. To illustrate the mathematical properties of the attention heads, we pictured separate "boxes" where each attention mechanism evolved in parallel, but in reality, they are slightly more connected. To fully utilize the efficient parallelization capability of the GPU hardware, it is critical to rethink every operation as a tensor operation. We described W^K, W^Q, and W^V of each head as separate matrices, but in practice, it is just three matrices that we conceptually break down by the number of heads needed.

Similarly, there is only one set of keys, queries, and values, and each head processes the entire sequence of tokens but operates on a distinct subset of features. The keys, queries, and values have dimension d_model X N, where N is the number of tokens in the input sequence. To specify each head's sub-segment explicitly, we reshape the matrices into 3-dimensional tensors with dimension n_head X d_head X N. Let us consider the incoming set of the hidden states. It is first projected into keys, queries, and values:

We then reshape the resulting matrices into 3-dimensional tensors:

The keys, queries, and values are reshaped into 3-dimensional tensors with dimension n_head X d_head X N, where each slice of the tensors corresponds to one head.

Reshaping is computationally efficient as it only reorganizes the tensor dimensions. When we compute the alignment scores E' from new tensors, this leads to N X N score for each head:

Here, we use the shorthand notation K’^T to streamline the notation and imply permutation on the last two indices of the tensor, similar to the transpose operation for matrices:

where k’_ijk is an element of K'. Notice that the way the operations are performed ensures the computation of N X N attention weights per head while keeping the number of arithmetic operations constant compared to the vanilla attention layer. The attention weights A' are obtained by normalizing on the last dimension:

again, e’_ijk is an element of the tensor E' and a’_ijk of the tensor A'. The context vectors are computed as the weighted average of the values with the attention weights:

or in tensor notation:

At this point, we have N context vectors of size d_head per head. We can reshape this tensor such that we have N context vectors of size d_model= n_head d_head:

We described this earlier as the concatenation of the different heads' context vectors. As a way to combine further the signal coming from the different heads, we pass the resulting context vectors through a final linear layer:

The computation of the attention and context vectors across multiple heads happens in parallel by making use of efficient tensor operations for GPU computing.

This approach lets the model process information more efficiently than sequential methods, making it better at understanding both nearby and far-apart relationships in the data.

The Positional Encoding

The Structure

The goal of the positional encoding (a.k.a. position embedding) in the Transformer architecture is to inject sequential order information into the model, enabling it to understand the position of tokens in a sequence. Since Transformers process all tokens in parallel (unlike sequential models like RNNs), they lack inherent awareness of token order. Position embeddings address this by encoding positional data. Without positional information, the Transformer would treat the input as a "bag of words," losing critical order-dependent structure.

In the "Attention is all you need" paper, the positional encoding is defined as another embedding matrix with the same embedding size as the token embedding. The number of rows in the position embedding defines the maximum number of tokens that the model can ingest within a sequence, also known as the context size. The positional information of the token is added to the model by summing the semantic vector representations of the tokens from the token embedding and their positional vector representations from the position embedding. This ensures that the self-attention weights carry the positional information such that the order of the tokens impacts the model inference.

The first set of hidden states is computed by summing the vector representations from the token embedding with the position encoding.

Without the position encoding, the model could not understand the order of the tokens in the sequence.

The position embedding (PE) is a static matrix of numbers. If i is the index position of the vectors in the embedding, and j is the index position of the elements in the vectors, the matrix elements are defined by the following formula:

where i ranges in [0, context size - 1] and j in [0, d_model - 1].

Capturing The Relative Token Positions

The motivation behind this choice of sinusoidal functional form for positional encodings is so the model can more easily learn attention weights reflecting each token's relative position. It stems from the trigonometric identities for sine and cosine functions:

Let us consider a fixed offset k, and we apply the trigonometric identities to the encoding formula:

where 𝛚_j = 1 / 1000^j/dmodel. In matrix notation, we have:

Let us define PE(i, j) as:

We obtain:

where:

In linear algebra, R_j(k) is called the rotation matrix and is used to perform a rotation in Euclidean space. Effectively, it means that PE(i+k, j) is the rotation of PE(i, j) by an angle -𝛚_jk.

So far, we have shown that, for two tokens with relative distance k, each pair of elements (j, j+1) within their positional encodings are related to each other through a rotation with angle -𝛚_jk. Let us call PE(i) = [PE(i, 0), PE(i, 1), …, PE(i, d_model)]. We can relate PE(i) and PE(i+k) through the pairwise rotation matrix:

where

Graphical representation of pairwise rotation of the different elements in the position encodings for two tokens separated by k tokens.

Let us now consider two hidden states h_i and h_i+k, corresponding to two tokens with relative distance k, coming into the self-attention layer. Both of them are the result of summing the token embedding vectors x_i and x_i+k and the positional encoding vectors PE(i) and PE(i+k):

We can compute their alignment score after projecting them into their keys and queries (we ignore heads for simplicity):

When we compute the alignment score between two tokens separated by k tokens, we can decompose its value into various contributions, including how the relative position k interacts with the absolute position i.

If we expend, we obtain:

We effectively decomposed the alignment score into four components:

Token-Token Interaction: Pure content-based alignment between x_i and x_i+k
Token-Position Interaction: How the token at i interacts with the relative position k of x_i+k
Position-Token Interaction: How the position i interacts with the token at x_i+k
Position-Position Interaction: How the relative position k (encoded via R(k)) interacts with the absolute position i.

Let's remember that R(k) is the fixed, mathematically defined transformation matrix (from sinusoidal identities) that maps PE(i) to PE(i+k) and it exists purely as a property of the positional encoding scheme. With this linear relationship, the model parameters W^K and W^Q can learn to leverage the structure of positional encodings to compute attention scores that depend on content and relative positions. During training, the model will learn to weigh these interactions by adjusting W^K and W^Q. It makes the training more efficient as the model does not need to relearn positional relationships from scratch; it builds on the mathematical structure of PE(i). The sinusoidal nature of the encoding also helps the model to generalize better to both unseen absolute positions and positional offsets.

Positional Encoding's Multi-Frequency Design

The positional encoding defines a frequency that depends on the position of vector elements:

Here, the constant 10,000 is a hyperparameter that controls the range of frequencies used to encode positional information. This means that the period of oscillations is 2𝝿 10000^j/dmodel. The frequencies range between [1, 1 / 10000^{(dmodel-1)/dmodel}]. High Frequencies lead to rapidly oscillating sine/cosine waves, which are adapted to distinguishing between nearby positions. This is crucial for local sentence syntax (e.g., word order in a phrase). Low frequencies lead to slowly oscillating waves that generalize over longer distances, which is useful for capturing global structure (e.g., paragraph-level coherence).

Each pair of columns captures different frequencies within the text data.

High Frequencies are adapted to distinguishing between nearby positions. This is crucial for local sentence syntax. Low frequencies are useful for capturing global structure.

The high value of 10,000 ensures a smooth transition from high to low frequencies across the embedding dimensions. 2𝝿 10000 is the maximum period supported by the model. Theoretically, this means it allows unique positional signals for tokens up to 2𝝿 10000 ~ 62,832 positions. However, transformers trained on sequences of fixed, shorter lengths (e.g., 512–4096 tokens) do not learn to handle positional relationships beyond this range. While the encoding theoretically supports very long periods, the model's effective context size is constrained by training data.

The Encoder

Introducing The Big Book of Large Language Models!

Damien Benveniste — Thu, 30 Jan 2025 16:01:31 GMT

For the past years, I have been creating educational content around machine learning and, specifically, large language models. I have been acquiring a depth of knowledge through my experience and practice in the field, and I want to share it with everybody! I started the process of writing, I believe, one of the most complete books on the subject of Large Language Models. You can access the book website here: The Big Book Of Large Language Models.

I will make the chapters available little by little as I write them. Don’t hesitate to leave comments so I can improve the current draft! The first chapter is now available: Language Models Before Transformers. In that chapter, I address the following subjects:

The Embedding Layers
Word2Vec
GloVe
The Jordan Network
The Elman Network
The Vanishing and Exploding Gradients Problem
Long Short Term Memory (LSTM)
Gated Recurrent Unit (GRU)
Sequence-to-Sequence Models
The RNN Encoder-Decoder Architecture
The Bahdanau Attention Mechanism
The Luong Attention

Here are the chapters coming up:

Introduction
Language Models Before Transformers
Attention Is All You Need: The Original Transformer Architecture
A More Modern Approach To The Transformer Architecture
Multi-modal Large Language Models
Transformers Beyond Language Models
Non-Transformer Language Models
How LLMs Generate Text
From Words To Tokens
Training LLMs to Follow Instructions
Scaling Model Training
Fine-Tuning LLMs
Deploying LLMs

My philosophy is to provide the depth of the mathematic notation along with the ease of visual illustrations of the different concepts. I believe the book can be read at different levels:

For somebody looking for the finest details, the equations should provide the foundations to understand thoroughly the concepts.
For somebody looking for a simpler read, the equation can be ignored to focus on the textual and visual explanations.
For somebody looking to strengthen their mathematical fundamentals in ML, the connection between the math and the visuals should help bridge the difficulties usually encountered when learning mathematics.

Let me know if you think the book is missing the target on that “mission.” I am truly excited to share this with you! I hope you will enjoy reading it as much as I enjoy writing it!