Into AI

A Little Intro to Context Management

To understand what went wrong, we need to understand how the LLM that acts as OpenClaw's brain manages its context.

An LLM uses its context window to store instructions and previous messages from an ongoing conversation. The context window functions as its short-term or working memory.

Although massive, these context windows are limited spaces that can fill up very quickly as the conversation grows.

Image from the author’s book titled ‘LLMs In 100 Images.’

As the context window approaches its end, response quality declines because the large amount of stored information makes it harder for the LLM to attend to it all.

An LLM tends to remember the beginning and end well, but starts forgetting the information buried in the middle of its context window. This shows up as hallucinations or the LLM ignoring instructions.

Once the context window is fully filled, the LLM stops answering user queries.

To ensure that such context window exhaustion is delayed as much as possible, previous conversation history goes through either:

Trimming: Where only a certain number of the most recent conversations are kept, and the previous ones are discarded before passing to an LLM at every turn.

Trimming conversations to keep only the last 2 recent ones and discarding the previous. (Other context for the LLM is omitted from the image for clarity.)

Summarization/ Compaction: Where the previous conversations are summarized and passed to the LLM at each turn, rather than the complete conversation history.

Summarise all conversations using the same LLM or another one, and give this summary to the LLM at each turn rather than the complete conversation history. (Other context for the LLM is omitted from the image for clarity.)

Pruning: Where old tool results (file reads, API responses, browser results) are selectively removed from the context while leaving all user queries and LLM responses untouched.

So What Really Went Wrong With OpenClaw?

Since Summer was previously using a small toy mailbox with OpenClaw, things worked pretty well. But when she moved the bot to use her actual mailbox, which was quite large, the context window filled up quickly, triggering Context compaction.

This led the bot to lose its original safety instruction on not taking action without confirmation, resulting in the deletion of hundreds of important emails.

Not really! It will still happen again if you don’t do the following.

10 Tips to Not End Up in the Same Place

If the director of Safety and Alignment at Meta Superintelligence Labs can make this mistake, this might happen to you as well. But there are ways to make sure that you do not end up in the same place. Here are 10 of them.

Put your security instructions in the MEMORY.md file rather than just in the conversation. This file is not compacted and acts as a persistent memory for the OpenClaw.
Then, you have the AGENTS.md file that is the standard operating and procedural rulebook that defines operational rules, security policies, and scope boundaries for OpenClaw to follow. Like MEMORY.md, it is loaded into context at the start of every session and survives compaction. Add an instruction here for the bot to check memory and get explicit confirmation before performing any destructive action.
OpenClaw can build a small vector index over MEMORY.md and memory/*.md so that semantic queries can find related notes even when the wording differs. Make sure that this is enabled in the ~/.openclaw/openclaw.json file.

{
  "agents": {
    "defaults": {
      "memorySearch": {
        "enabled": true,
        "provider": "local"
      }
    }
  }
}

OpenClaw has a mechanism called Pre-compaction memory flush that is automatically triggered in the background as the context window is about to fill up. This saves all the important details from context to disk before compaction starts.

To enable this, go to the ~/.openclaw/openclaw.json file and make sure it contains the following:

{
  "agents": {
    "defaults": {
      "compaction": {
        "reserveTokensFloor": 20000, // Tokens reserved as buffer before compaction triggers
        "memoryFlush": {
          "enabled": true, 
          "softThresholdTokens": 4000, // // How far before the reserve floor the flush triggers
          "systemPrompt": "Session nearing compaction. Store durable memories now.", 
          "prompt": "Write any lasting notes to memory/YYYY-MM-DD.md; reply with NO_REPLY if nothing to store."  
        }
      }
    }
  }
}

Enable Context pruning, which removes old tool results from the context before sending to the bot. Add the following to the ~/.openclaw/openclaw.json file:

{
  “agents”: {
    “defaults”: {
      “contextPruning”: {
        “mode”: “cache-ttl”, // Automatically remove tool results after TTL
        “ttl”: “6h” // Keep tool results for 6 hours
      }
    }
  }
}

Use the /context list or /context detail commands in chat regularly to learn what is in the bot's context window and whether any contextual files are truncated or missing.
Always save important instructions manually to MEMORY.md file by typing “Save this to MEMORY.md“ in your chat with OpenClaw.
Don’t wait for compaction to happen automatically. Instead, perform it manually. Use the /compact command in chat when sessions feel stale, or context is getting bloated, and then give your new instructions to the bot.
If you still end up with OpenClaw misbehaving, sending /stop in the chat aborts the session and stops any active sub-agent runs spawned from it, including the nested children.
Last but not least, if none of the above works, rush to the machine running OpenClaw and pull the power plug. This should most definitely stop it.

A very helpful lesson on OpenClaw memory can be found here if you’re interested in learning about it in more detail.

P.S. This post isn't meant to mock Summer Yue, but it’s quite the opposite. She is one of the very few people who openly admitted a mistake in public that most would hide, given her job role. Thanks to her, we're all better engineers because she chose transparency over ego.

Join the paid tier today to get access to all posts on this newsletter:

and so many more!

You can also read my books on Gumroad and connect with me on LinkedIn to stay in touch.

🗓️ This Week In AI Research (15-21 March 26)

Dr. Ashish Bamania — Wed, 25 Mar 2026 10:27:50 GMT

1. Mixture-of-Depths Attention

As LLMs become deeper, they often suffer from signal degradation. This means that important features formed in shallow layers are gradually diluted by repeated residual updates, making them harder to recover in deeper layers.

This research paper introduces the Mixture-of-depths attention (MoDA), which allows each attention head to attend to sequence KV pairs at the current layer and to depth KV pairs from preceding layers.

The authors also describe a hardware-efficient algorithm for MoDA that resolves non-contiguous memory-access patterns, achieving 97.3% of FlashAttention-2’s efficiency at a sequence length of 64K.

Experiments on 1.5B-parameter models show that MoDA consistently outperforms strong baselines, and combining MoDA with post-norm yields better performance than using it with pre-norm.

Read more about this research paper using this link.

🎁 Join the paid tier today to get access to all posts on this newsletter.

and so many more!

2. Attention Residuals

This research paper proposes Attention Residuals (AttnRes), an alternative to standard residual connections in transformers.

Typically, models add all outputs from previous layers equally, which can dilute valuable information as the model deepens. AttnRes instead allows each layer to use attention to selectively merge earlier layers, assigning learned, input-dependent weights to focus on the most relevant past information.

To increase efficiency, Block AttnRes groups layers into blocks and attends to summaries of these blocks rather than every layer, reducing memory costs while preserving most benefits of full AttnRes.

Overall, this enhances training stability, maintains stronger signals across depth, and improves performance, particularly in large models.

Read more about this research paper using this link.

3. Mamba-3

This research paper introduces Mamba-3, a new sequence model that makes inference much more efficient than transformers (which have quadratic compute and high memory costs).

The key idea is an inference-first redesign of state space models (SSMs) where Mamba-3 uses a:

More expressive recurrence
Complex-valued state update to better track information over time
Multi-input multi-output (MIMO) setup that improves performance without slowing down decoding

These enable it to handle difficult tasks such as state tracking and long-range dependencies, areas where previous linear models struggled, while maintaining near-linear efficiency. This results in Mamba-3 also achieving better accuracy and similar or lower memory use compared to previous models.

Read more about this research paper using this link.

4. V-JEPA 2.1

This research paper presents V-JEPA 2.1, a family of self-supervised models that learn dense, high-quality visual representations for both images and videos while retaining strong global scene understanding.

The approach combines four key components.

A dense predictive loss that uses a masking-based objective in which both visible and masked tokens contribute to the training signal, encouraging explicit spatial and temporal grounding.
Deep self-supervision that applies the self-supervised objective hierarchically across multiple intermediate encoder layers to improve representation quality.
Multi-modal tokenizers that enable unified training across images and videos.
Effective scaling in both model capacity and training data

Empirically, V-JEPA 2.1 achieves state-of-the-art performance on several challenging benchmarks for short-term object-interaction anticipation and high-level action anticipation.

The model also shows strong performance in robotic navigation, depth estimation, and global recognition.

Read more about this research paper using this link.

5. GPT‑5.4 mini and nano

OpenAI released GPT‑5.4 mini and nano, which are two small, fast, and economical versions of GPT-5.4.

GPT-5.4 mini is a strong, lightweight general-purpose model, especially for coding, reasoning, multimodal understanding, and tool use, and runs over 2× faster than GPT-5-mini. It also approaches the performance of the larger GPT‑5.4 model on several evaluations, including SWE-Bench Pro and OSWorld-Verified.

GPT‑5.4 nano is the smallest, cheapest version of GPT‑5.4 for tasks where speed and cost matter most, such as classification, data extraction, ranking, and coding subagents that handle simpler supporting tasks.

Read more about this release using this link.

6. Composer-2

Composer 2 is a frontier-level coding model that offers a new, optimal combination of intelligence and cost, delivering impressive performance on hard benchmarks, including Terminal-Bench 2.0 and SWE-bench Multilingual.

The model is the result of a continued pretraining run (built on top of the open-source Kimi K2.5 model from Moonshot AI) and further training on long-horizon coding tasks through reinforcement learning, enabling it to solve challenging tasks requiring hundreds of actions.

Read more about this release using this link.

A technical report can be found using this link.

7. BioReason-Pro

This research paper introduces BioReason-Pro, a new multimodal LLM for predicting protein functions that goes beyond traditional methods that rely on sequence similarity or classification.

Instead, it explicitly models biological reasoning across sequence, structure, domains, and interaction data.

BioReason-Pro combines a new autoregressive Gene Ontology predictor (GO-GPT) with a reasoning LLM trained on large-scale synthetic reasoning traces and further improved using reinforcement learning, enabling it to generate clear, step-by-step functional annotations.

Empirically, BioReason-Pro achieves strong performance and produces higher-quality functional summaries than previous methods. Human experts prefer its annotations over curated UniProt entries about 79% of the time.

Notably, the model shows emergent capabilities such as:

accurately predicting protein binding partners
identifying interaction sites at the residue level, consistent with experimental structures
correcting misleading annotations by using contextual reasoning

Read more about this research using this link.

8. Hyperagents

This research paper introduces Hyperagents, a framework that combines a task agent, which solves problems, and a meta agent, which improves the system, into a single editable program.

Importantly, the improvement process is also editable, allowing metacognitive self-modification, meaning the system enhances its own improvement.

The research builds on the Darwin Gödel Machine (DGM) to create DGM-Hyperagents (DGM-H), which removes DGM's assumption that self-improvement only works within coding and outperforms baselines across coding, paper review, robotics reward design, and Olympiad-level math-solution grading, showing meta-level improvements that transfer across domains.

Read more about this research using this link.

9. LeWorldModel

This research paper introduces LeWorldModel (LeWM), which is the first Joint Embedding Predictive Architecture (JEPA) that trains stably end-to-end from raw pixels using only two loss terms:

a next-embedding prediction loss
a regularizer enforcing Gaussian-distributed latent embeddings

This reduces the number of tunable loss hyperparameters from 6 to 1 compared to the only existing end-to-end alternative.

With approximately 15M parameters trainable on a single GPU in a few hours, LeWM plans up to 48x faster than foundation-model-based world models while remaining competitive across diverse 2D and 3D control tasks.

Beyond control, the researchers show that LeWM's latent space encodes meaningful physical structure through probing of physical quantities. Surprise evaluation confirms that the model reliably detects physically implausible events.

Read more about this research using this link.

10. Nemotron-Cascade 2

This research paper from NVIDIA introduces Nemotron-Cascade 2, an open 30B MoE model with 3B activated parameters that delivers best-in-class reasoning and strong agentic capabilities.

Despite its compact size, its mathematical and coding reasoning performance approaches that of frontier open models. It is the second open-weight LLM, after DeepSeekV3.2-Speciale-671B-A37B, to achieve Gold Medal-level performance in the 2025 International Mathematical Olympiad (IMO), the International Olympiad in Informatics (IOI), and the ICPC World Finals, demonstrating remarkably high intelligence density with 20x fewer parameters.

Read more about this research using this link.

This article is entirely free to read. If you loved reading it, restack and share it with others to earn referral rewards. ❤️

Join the paid tier today to get access to all posts on this newsletter:

and so many more!

Link to the Visual Tech Bundle

You can also read my books on Gumroad and connect with me on LinkedIn to stay in touch.

Build and Train a Mixture-of-Experts (MoE) LLM from Scratch

Dr. Ashish Bamania — Fri, 20 Mar 2026 11:51:20 GMT

Most modern-day LLMs use the Mixture of Experts (MoE) architecture. This includes Grok-1, DeepSeekMoE, gpt-oss, and Mixtral (and many other proprietary LLMs whose architectural details aren’t publicly available).

In the previous lessons on Into AI, we learned:

In this lesson, we will:

Build an LLM using the MoE transformer with Grouped Query Attention (GQA)
Train it on a publicly available database from Wikipedia
Generate text from the trained LLM

The following are the steps that we implement.

Let’s begin!

Import Necessary Packages

# Standard libraries
import math
import random

# Hide deprecation warnings
import warnings
warnings.filterwarnings('ignore')

# PyTorch core
import torch
import torch.nn as nn
import torch.optim as optim

# For processing data
from datasets import load_dataset
from torch.utils.data import Dataset, DataLoader

# BPE tokenizer
import tiktoken

# Training utils
from torch.nn.utils import clip_grad_norm_
from torch.optim.lr_scheduler import CosineAnnealingLR

# Progress bar
from tqdm import tqdm

Before we move forward, I want to introduce you to the Visual Tech Bundle.

It is a collection of visual guides that explain core AI, LLM, Systems design, and Computer science concepts via image-first lessons.

Others are already loving these books.

This includes Dharmesh Shah, the co-founder and CEO of HubSpot.

Dharmesh’s reply to my LinkedIn post. (Source)

❤️ I’m sure that you will love these too:

Now back to our lesson!

Getting our data ready

1. Download the dataset

We will use the WikiText dataset from Hugging Face to train our model. This dataset is derived from verified Wikipedia articles and contains approximately 103 million words.

It is downloaded as follows.

# Load WikiText dataset 
train_data = load_dataset("wikitext", "wikitext-103-v1", split="train") # Train subset
val_data   = load_dataset("wikitext", "wikitext-103-v1", split="validation") # Validation subset

Let’s check out a training example from this dataset.

# Training subset example
print(train_data['text'][807])

"""
Output:
The International Civil Aviation Organization ( ICAO ) defines general aviation ( GA ) 
as " an aircraft operation other than a commercial air transport operation 
or an aerial work operation . " It defines commercial air transport ( CAT ) 
as " an aircraft operation involving the transport of passengers , cargo 
or mail for remuneration or hire " , and aerial work as " an aircraft operation 
in which an aircraft is used for specialized services such as agriculture , 
construction , photography , surveying , observation and patrol , search and 
rescue , aerial advertisement , etc . " 

"""

There’s a little problem with this dataset where some text examples are blanks and headings are formatted as “ === HEADING === “, as shown below.

# Unwanted training subset example
print(train_data['text'][804:807])

"""
Output:
['', ' = = Definitions = = \n', '']
"""

Alongside this, WikiText-103 was built with a vocabulary consisting of words (rather than sub-words/ characters), and any word that appeared fewer than three times was replaced with .

This is shown below.

# UNK marker
print(train_data['text'][80])

"""
Output:
   2 
"""

We need to get rid of these to clean up the dataset.

We also need an End-of-Text (EOS) token called <|endoftext|> between separate sections and articles, so that the model learns natural stopping boundaries and doesn't try to connect unrelated content.

(We will come back to re-using the EOS token when discussing the helper function for text generation/inference.)

Both of these changes are made using a helper function as follows.

def clean_text(dataset):
    # EOS token
    EOS = "<|endoftext|>"

    cleaned = []
    for text in dataset["text"]:
        # Strip surrounding whitespace
        line = text.strip()

        # Keep only non-empty lines that aren't headings (such as "=== HEADING ===")
        if line and not line.startswith("="):
            # Remove  markers and normalize whitespace
            cleaned.append(" ".join(line.replace("", "").split()))

    # Join all entries into a string, separated by EOS token
    return EOS.join(cleaned) + EOS
 
# Clean training and validation text
training_text = clean_text(train_data)
validation_text = clean_text(val_data)

A short subset of training_text is shown below.

# Print a subset of 'training_text'
print(training_text[3457:3700])

"""
Output:
Each character has a field and distance of movement limited by their 
Action Gauge . Up to nine characters can be assigned to a single mission . 
During gameplay , characters will call out if something happens to them , 
such as their health points ( HP ) getting low or being knocked out
"""

2. Tokenize the dataset

We previously learned how to build a character-level tokenizer from scratch.

We have also used it to train an LLM from scratch.

But because it is character-based, training an LLM using it makes the model learn the meaning of text from individual letters rather than semantic units. This makes training painfully difficult.

For this tutorial, we will instead use Tiktoken, a fast BPE tokeniser used in OpenAI's models. It is a sub-word tokenizer (breaks down words into sub-words) that better captures language structure than a character-level tokenizer, enabling faster LLM training.

Image from the author’s book titled ‘LLMs In 100 Images.’

We start by creating an instance of Tiktoken as follows.

# Create an instance of Tokenizer used by GPT-2
tokenizer = tiktoken.get_encoding("gpt2")

The vocabulary size when using the Tiktonen tokenizer for GPT-2 is 50,257. This means that the tokenizer recognizes 50,257 unique subwords.

# Check vocabulary size
vocab_size = tokenizer.n_vocab

print(f"Vocabulary size: {vocab_size}")

# Output: Vocabulary size: 50257

3. Create a dataset required for language modeling

Once we have tokenized our dataset, we need to load and serve the data during training. This is where the TextDataset class comes in, which inherits from PyTorch’s Dataset class, and has the following methods:

__init__: Tokenizes given text (tokens to token IDs) and sets the maximum sequence length (max_seq_length)
__len__: Returns the number of training sequences available
__getitem__: Returns a training sequence and its targets (tokens shifted forward by one position) at a given index

# Dataset for language modeling
class TextDataset(Dataset):
    def __init__(self, text, tokenizer, max_seq_length):
        # Convert text to token IDs
        self.tokens = tokenizer.encode(text, allowed_special="all")

        # Maximum length of each training sequence
        self.max_seq_length = max_seq_length

    # Get number of valid training sequences
    def __len__(self):
        num_sequences = (len(self.tokens) - 1) // self.max_seq_length
        return num_sequences

    # Get an input sequence and targets
    def __getitem__(self, idx):
        # Start index of the sequence
        start = idx * self.max_seq_length

        # End index of the sequence
        end = start + self.max_seq_length

        # Input token sequence
        input_ids = torch.tensor(self.tokens[start:end], dtype=torch.long)

        # Next-token targets/ labels (shifted by one character)
        target_ids = torch.tensor(self.tokens[start+1:end+1], dtype=torch.long)

        return input_ids, target_ids

Let’s use the TextDataset class to create the training and validation datasets.

# Define maximum sequence length
MAX_SEQ_LENGTH = 128

# Create training and validation datasets
train_dataset = TextDataset(training_text, tokenizer, MAX_SEQ_LENGTH)
val_dataset = TextDataset(validation_text, tokenizer, MAX_SEQ_LENGTH)

The number of training sequences in the training and validation datasets is as follows.

print(f"Number of training sequences: {len(train_dataset):,}")
print(f"Number of validation sequences: {len(val_dataset):,}")

"""
Output:
Number of training sequences: 876,162
Number of validation sequences: 1,828
"""

And this is how a training sequence and its target look.

# Example of a training sequence and its target
input, target = train_dataset[16]

print("Input IDs:\n", input)
print("\nTarget IDs:\n", target)

print("\nDecoded Input:\n", tokenizer.decode(input.tolist()))
print("\nDecoded Target:\n", tokenizer.decode(target.tolist()))

"""
Output:
Input IDs:
 tensor([  262,  1271,   286, 22469,  4991,   583,  4365,   764,   317,   636,
          286,   428,  8515,  2950,  4441,  3748,  7514, 14520,  4981,   329,
         1123,  2095,   705,    82,  1767,   764,   554,  1502,   284,  4620,
          428,   837,   262, 22849,  4847, 16560,   656,   262,  1218,   983,
          547,  4615,   837,   355,   484,  1718,   510,   257,  1588,  6903,
          286,  4088,  2272,  2622,   329,   262,  8561,   764,  1119,   635,
        12328,   262,  8722,  6460,   290, 10152,   286,   711,   523,   484,
          714,  5198,   284,   649,  1938,   981, 26645,   262,  6393,  6805,
          286,   262,  2168,   705, 11327,   764,   383, 15064,  3341,   547,
         3066,  2402,  1903,   287,  2478,   764,   383,  2095,  9824,   547,
         1760,   416,  8835,    73,   280,   837,   508,   550,  3111,   319,
          262,  2180,   569, 18354,  7496, 17740,  1830,   764,  1649,  4441,
          262, 17871,  5321, 11630,   837,  8835,    73,   280])

Target IDs:
 tensor([ 1271,   286, 22469,  4991,   583,  4365,   764,   317,   636,   286,
          428,  8515,  2950,  4441,  3748,  7514, 14520,  4981,   329,  1123,
         2095,   705,    82,  1767,   764,   554,  1502,   284,  4620,   428,
          837,   262, 22849,  4847, 16560,   656,   262,  1218,   983,   547,
         4615,   837,   355,   484,  1718,   510,   257,  1588,  6903,   286,
         4088,  2272,  2622,   329,   262,  8561,   764,  1119,   635, 12328,
          262,  8722,  6460,   290, 10152,   286,   711,   523,   484,   714,
         5198,   284,   649,  1938,   981, 26645,   262,  6393,  6805,   286,
          262,  2168,   705, 11327,   764,   383, 15064,  3341,   547,  3066,
         2402,  1903,   287,  2478,   764,   383,  2095,  9824,   547,  1760,
          416,  8835,    73,   280,   837,   508,   550,  3111,   319,   262,
         2180,   569, 18354,  7496, 17740,  1830,   764,  1649,  4441,   262,
        17871,  5321, 11630,   837,  8835,    73,   280,   373])

Decoded Input:
  the number of playable units per mission . A part of this upgrade involved creating 
  unique polygon models for each character 's body . In order to achieve this , 
  the cooperative elements incorporated into the second game were removed , 
  as they took up a large portion of memory space needed for the improvements . 
  They also adjusted the difficulty settings and ease of play so they could appeal 
  to new players while retaining the essential components of the series ' gameplay . 
  The newer systems were decided upon early in development . The character designs 
  were done by Honjou , who had worked on the previous Valkyria Chronicles games . 
  When creating the Nameless Squad , Honjou

Decoded Target:
  number of playable units per mission . A part of this upgrade involved creating 
  unique polygon models for each character 's body . In order to achieve this , 
  the cooperative elements incorporated into the second game were removed , 
  as they took up a large portion of memory space needed for the improvements . 
  They also adjusted the difficulty settings and ease of play so they could appeal 
  to new players while retaining the essential components of the series ' gameplay . 
  The newer systems were decided upon early in development . The character designs 
  were done by Honjou , who had worked on the previous Valkyria Chronicles games . 
  When creating the Nameless Squad , Honjou was
"""

Note how the input and target sequences are shifted by one sub-word. This is to help the model learn to predict each next sub-word given the previous ones.

Setting up DataLoader for Batch training

Next, we use PyTorch DataLoader to create an iterable that loads shuffled batches of BATCH_SIZE samples from train_dataset and val_dataset for training and validation.

# Define batch size
BATCH_SIZE = 32

# Create training and validation DataLoaders
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

This creates the following number of batches of training and validation data.

print(f"Number of training batches: {len(train_loader):,}") # Number of training sequences / batch_size
print(f"Number of validation batches: {len(val_loader):,}") # Number of validation sequences / batch_size

"""
Output:
Number of training batches: 27,381
Number of validation batches: 58
"""

Our data is now ready for model training!

Building the MoE Transformer model

We will now set up our MoE transformer model that we built in a previous lesson with a few changes (based on what modern LLMs follow):

We replace the Causal Multi-head Self-attention module with Causal Grouped-Query Attention
We introduce Weight tying, which means we reuse the same weights for both the token embedding matrix and the output projection matrix, rather than learning two separate ones. This is because they do reciprocal functions and have the same shape. This decreases the number of learning parameters in the model by a few million.
We remove the bias term from all linear layers following the Llama models

Let’s implement these one by one.

1. Set up Causal Grouped Query Attention

# Grouped Query Attention (GQA) with Causal Masking
class GroupedQueryAttention(nn.Module):
    def __init__(self, embedding_dim, num_heads, num_groups, max_seq_length):
        super().__init__()
 
        # Check if embedding_dim is divisible by num_heads
        assert embedding_dim % num_heads == 0, "embedding_dim must be divisible by num_heads"
 
        # Check if num_heads is divisible by num_groups
        # (Each group must be shared by the same number of heads)
        assert num_heads % num_groups == 0, "num_heads must be divisible by num_groups"
 
        # Embedding dimension
        self.embedding_dim = embedding_dim
 
        # Number of total query heads
        self.num_heads = num_heads
 
        # Dimension of each head
        self.head_dim = embedding_dim // num_heads
 
        # Number of KV groups
        self.num_groups = num_groups
 
        # Number of query heads per KV group
        self.group_size = num_heads // num_groups
 
        # Linear projection matrix for query
        self.W_q = nn.Linear(embedding_dim, embedding_dim, bias=False)
 
        # Linear projection matrices for key and value
        self.W_k = nn.Linear(embedding_dim, num_groups * self.head_dim, bias=False)
        self.W_v = nn.Linear(embedding_dim, num_groups * self.head_dim, bias=False)
 
        # Linear projection matrix to produce final output
        self.W_o = nn.Linear(embedding_dim, embedding_dim, bias=False)
 
        # Build the causal mask once at init instead of every forward pass
        mask = torch.tril(torch.ones(max_seq_length, max_seq_length))
 
        # Add batch_size and num_heads dimensions
        mask = mask.view(1, 1, max_seq_length, max_seq_length)

        # register_buffer saves it as part of the model so it automatically moves to the right device (CPU/GPU)
        self.register_buffer("causal_mask", mask)
 
    # Splits Q into multiple heads
    def _split_heads(self, x):
        """
        Transforms input embeddings from
        [batch_size, sequence_length, embedding_dim]
        to
        [batch_size, num_heads, sequence_length, head_dim]
        """
        batch_size, sequence_length, embedding_dim = x.shape
 
        # Split embedding_dim into (num_heads, head_dim)
        x = x.reshape(batch_size, sequence_length, self.num_heads, self.head_dim)
 
        # Reorder and return the intended shape
        return x.transpose(1, 2)
 
    # Splits K or V into num_groups heads
    def _split_groups(self, x):
        """
        Transforms K/V from
        [batch_size, sequence_length, num_groups * head_dim]
        to
        [batch_size, num_groups, sequence_length, head_dim]
        """
        batch_size, sequence_length, _ = x.shape
 
        x = x.reshape(batch_size, sequence_length, self.num_groups, self.head_dim)
        return x.transpose(1, 2)
 
    # Merge heads back together
    def _merge_heads(self, x):
        """
        Transforms inputs from
        [batch_size, num_heads, sequence_length, head_dim]
        to
        [batch_size, sequence_length, embedding_dim]
        """
        batch_size, num_heads, sequence_length, head_dim = x.shape
 
        # Move sequence_length back before num_heads in the shape
        x = x.transpose(1, 2)
 
        # Merge (num_heads, head_dim) back into embedding_dim
        embedding_dim = num_heads * head_dim
        x = x.reshape(batch_size, sequence_length, embedding_dim)
 
        return x
 
    # Forward pass
    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.shape
 
        # Compute Q, K, V
        Q = self.W_q(x)  # [batch_size, sequence_length, embedding_dim]
        K = self.W_k(x)  # [batch_size, sequence_length, num_groups * head_dim]
        V = self.W_v(x)  # [batch_size, sequence_length, num_groups * head_dim]
 
        # Split Q into multiple heads
        Q = self._split_heads(Q)  # [batch_size, num_heads, sequence_length, head_dim]
 
        # Split K and V into num_groups heads
        K = self._split_groups(K)  # [batch_size, num_groups, sequence_length, head_dim]
        V = self._split_groups(V)  # [batch_size, num_groups, sequence_length, head_dim]
 
        # Expand K and V so each KV group is shared across multiple query heads
        K = K.repeat_interleave(self.group_size, dim=1)  # [batch_size, num_heads, sequence_length, head_dim]
        V = V.repeat_interleave(self.group_size, dim=1)  # [batch_size, num_heads, sequence_length, head_dim]
 
        # Calculate scaled dot-product attention
        attn_scores = Q @ K.transpose(-2, -1)
        attn_scores = attn_scores / math.sqrt(self.head_dim)
 
        # Slice the pre-built mask to the current sequence length
        causal_mask = self.causal_mask[:, :, :sequence_length, :sequence_length]
 
        # Mask out future positions by setting their scores to -inf
        attn_scores = attn_scores.masked_fill(causal_mask == 0, float('-inf'))
 
        # Apply softmax to get attention weights
        attn_weights = torch.softmax(attn_scores, dim=-1)
 
        # Multiply attention weights by V
        weighted_values = attn_weights @ V  # [batch_size, num_heads, sequence_length, head_dim]
 
        # Merge head outputs
        merged_heads_output = self._merge_heads(weighted_values)
 
        # Obtain final output
        output = self.W_o(merged_heads_output)
 
        return output

2. Set up Expert and Router

An Expert is simply a small feed-forward network that specializes in a particular skill, which may involve processing different token types (numerical/code/language-related tokens), reasoning patterns, context patterns, and so on.

🗓️ This Week In AI Research (8-14 March 26)

Dr. Ashish Bamania — Thu, 19 Mar 2026 11:14:13 GMT

1. Covenant-72B

This research paper presents Covenant-72B, an LLM produced by the largest globally distributed collaborative pre-training run to date, in terms of both compute and model scale.

This pre-training involved open, permissionless participation, supported by a live blockchain protocol, Bittensor, and a state-of-the-art, communication-efficient optimizer called SparseLoCo, which enabled dynamic participation with peers joining and leaving freely.

Covenant-72B was pre-trained on approximately 1.1T tokens and performs at par with fully centralized models pre-trained on similar or higher compute budgets.

This shows that open and fully democratized participation is possible on an unprecedented scale for a globally distributed LLM pre-training run.

Read more about this research using this link.

🎁 Unlock every lesson in the newsletter and accelerate your AI engineering career with a 25% discount on the annual membership of ‘Into AI’ today!

2. Training Language Models via Neural Cellular Automata

Natural LLM pre-training data has its flaws, as it is finite, contains human biases, and mixes knowledge with reasoning.

This research paper proposes using neural cellular automata (NCA) to generate synthetic, non-linguistic data for pre-pre-training LLMs, which means training on synthetic data first, followed by natural language.

This is because NCA data has a rich spatiotemporal structure and statistics that resemble natural language, while being controllable and cheap to generate at scale.

Pre-pre-training is surprisingly effective, and training on only 164M NCA tokens improves downstream language modeling by up to 6% and accelerates convergence by up to 1.6x.

This approach of using NCA even outperforms pre-pre-training on 1.6B tokens of natural language from Common Crawl with more compute, and the gains transfer to reasoning benchmarks, such as GSM8K, HumanEval, and BigBench-Lite.

Further experiments show that attention layers transfer best, and that optimal NCA complexity depends on the task, with simpler dynamics working better for code, while more complex ones benefit math and web text more.

Read more about this research using this link.

3. Lost in Backpropagation: The LM Head is a Gradient Bottleneck

The final layer of LLMs, called the LM head, maps the hidden dimension (D) to the vocabulary dimension (V), where usually D is very small as compared to V. This creates an expressivity limitation called the Softmax bottleneck.

This research paper shows that the softmax bottleneck is not only an expressivity bottleneck but also an optimization bottleneck where high-dimensional gradients are compressed into a much smaller space, leading to most of the training signal (around 95–99% of the gradient norm) getting lost.

This leads to poor update directions and makes even simple patterns harder to learn, slowing down training and harming efficiency across models.

The research suggests that current LLM architectures have a fundamental flaw in how they learn and that better LM head designs are needed.

Read more about this research using this link.

4. OpenClaw-RL: Train Any Agent Simply by Talking

This research paper presents OpenClaw-RL, a reinforcement learning framework in which AI agents learn continuously from real-time interactions, using next-state signals (such as user replies, tool outputs, or environment changes) as training data.

These signals offer evaluative feedback (how well an action is performed, expressed as scalar rewards via a PRM judge) and directive feedback (how the action should be improved, expressed as textual hints and used for token-level supervision via Hindsight-Guided On-Policy Distillation, or OPD).

By blending these, OpenClaw-RL supports more effective learning than traditional methods.

Its asynchronous, decoupled architecture also allows agents to improve during live use in various environments, such as chat, tools, and software tasks, without requiring offline datasets or manual labeling.

Read more about this research using this link.

5. How Far Can Unsupervised RLVR Scale LLM Training?

This research paper studies a training approach called Unsupervised Reinforcement Learning with Verifiable Rewards (RLVR), in which models learn without labeled data by generating their own rewards from signals that can be automatically verified.

The paper categorises RLVR methods into intrinsic (self-generated rewards) and external (externally verified rewards), and shows that intrinsic methods mainly strengthen a model’s existing beliefs and work only if those beliefs are already correct, but can lead to collapse if they are wrong.

While intrinsic rewards still help with test-time training on small datasets, the authors propose the Model Collapse Step, which acts as a practical indicator for RL trainability (predicting when training will fail).

They also present early evidence that external reward methods may overcome these limitations, suggesting more scalable approaches.

Read more about this research using this link.

6. Attend Before Attention

This research paper introduces AutoGaze for multi-modal LLMs (MLLMs), an efficient method for processing videos by selectively focusing on important regions (like human gaze) before applying full attention, rather than analyzing every pixel equally.

Such gazing lets the model ignore unnecessary or static information, enabling faster, easier-to-scale video understanding while maintaining high performance.

Experiments show that this approach greatly increases efficiency (higher processing FPS) and allows for real-time video processing, offering a practical way to scale multimodal models to longer or streaming video inputs.

Read more about this research using this link.

7. Geometry-Grounded Point Transformer

Feed-forward networks suffer from geometric inconsistencies and achieve only limited fine-grained accuracy in 3D reconstruction (i.e., predicting dense point maps directly from RGB images).

This research paper introduces the Geometry-Grounded Point Transformer (GGPT), a method for improving 3D reconstruction from images by combining fast neural predictions with explicit geometric guidance.

It uses an improved Structure-from-Motion pipeline to estimate accurate camera poses and partial 3D structure from sparse input views, and then uses a Geometry-guided 3D point transformer to refine dense point maps under explicit partial-geometry supervision using an optimised guidance encoding.

Experiments show that GGPT produces reconstructions that are both geometrically consistent and spatially complete.

While being trained solely on ScanNet++ with VGGT predictions, it generalises across architectures and datasets, significantly outperforming SOTA feed-forward 3D reconstruction models in both in-domain and out-of-domain settings.

Read more about this research using this link.

8. The Attack and Defense Landscape of Agentic AI

This research paper is a comprehensive survey of security risks in AI agent systems that examines how design choices, such as tool use, memory, and autonomy, create new attack surfaces.

It uses case studies to highlight existing gaps in securing agentic AI systems and identifies open challenges in this domain.

It also introduces the first systematic framework for understanding the security risks and defense strategies of AI agents, serving as a foundation for building secure agentic systems and advancing research.

Read more about this research using this link.

9.IndexCache

This research paper presents IndexCache, a method that improves sparse attention in LLMs by reusing attention indices across layers rather than recalculating them each time.

In a standard sparse attention mechanism, models repeatedly select the most relevant tokens at each layer, which is computationally expensive. This research shows that these selections are often similar across layers and can be cached and reused.

By adding this cross-layer index reuse method, the approach reduces computation and latency while maintaining model performance, making long-context and agent-style workloads more efficient and scalable.

Read more about this research using this link.

10. Neural Thickets

This research paper suggests that large pre-trained LLMs already have many good solutions for task-specific experts nearby in their weights. They do not require extensive training or optimization to uncover them.

Instead of using complex optimization methods, the authors propose a simple, fully parallel post-training approach where they make small random changes to the model parameters, select the top-performing variants, and combine their predictions using majority voting.

This straightforward approach is surprisingly effective and works almost as well as standard post-training methods such as PPO, GRPO, and ES for contemporary large-scale models.

Read more about this research using this link.

This article is entirely free to read. If you loved reading it, restack and share it with others to earn referral rewards. ❤️

Join the paid tier today at a 25% discount to get access to all posts on this newsletter!

You can also check out my books on Gumroad and connect with me on LinkedIn to stay in touch.

🗓️ This Week In AI Research (1-7 March 26)

Dr. Ashish Bamania — Tue, 10 Mar 2026 12:38:02 GMT

1. Introducing GPT‑5.4

OpenAI released GPT-5.4, with improved reasoning, coding, and professional knowledge capabilities, enabling the development of more dependable AI agents.

GPT-5.4 can perform hard tasks across documents, spreadsheets, presentations, and software tools, producing higher-quality outputs and helping teams complete workflows faster.
It is the first general-purpose model with native, state-of-the-art computer-use capabilities.
It supports up to 1M tokens of context, allowing agents to plan, execute, and verify tasks across long horizons.

It is available across ChatGPT, the OpenAI API, and Codex, with variants such as GPT-5.4 Thinking for deeper reasoning and GPT-5.4 Pro for more demanding workloads.

Read more about this release using this link.

🚨 A Special Offer

Grab a 25% discount on the annual subscription to the newsletter and get access to articles like these, and so much more:

Your 25% off

2. FlashAttention-4

This research paper introduces the latest version, FlashAttention-4, which makes transformer attention faster and more efficient on modern GPUs.

FlashAttention-3 was optimized attention for Hopper GPUs, specifically targeting the H100 architecture. But the industry has moved to Blackwell-based systems such as the B200 and GB200, which come with fundamentally different performance characteristics.

FlashAttention-4 optimizes attention on Blackwell GPUs by:

Redesigning pipelines that take advantage of fully asynchronous Matrix Multiply-Accumulate (MMA) operations and larger tile sizes
Software-based implementations of exponential and conditional softmax rescaling, which reduce the number of non-matrix operations and improve overall speed
Better use of GPU tensor memory and the 2-CTA MMA mode (which allows two Cooperative Thread Arrays to jointly execute a single Matrix Multiply-Accumulate operation), which reduces shared-memory traffic and atomic adds required during the backward pass

FlashAttention-4 achieves up to:

a 1.3× speedup over cuDNN 9.13
a 2.7× speedup over Triton on B200 GPUs with BF16, reaching up to 1613 TFLOPs/s (71% utilization)

The researchers also implement FlashAttention-4 entirely in CuTe-DSL embedded in Python, resulting in 20-30× faster compile times compared to traditional C++ template-based approaches while maintaining full expressivity.

Read more about this research using this link.

3. Solving an Open Problem in Theoretical Physics using AI-Assisted Discovery

This research paper describes a neuro-symbolic system that autonomously solved an open problem in theoretical physics.

This system combines Gemini Deep Think with a systematic Tree Search (TS) algorithm and an automated numerical feedback system to explore mathematical ideas and verify them computationally.

The system successfully derived new and exact analytical solutions for the power spectrum of gravitational radiation emitted by cosmic strings.

The research system identified 6 different analytical methods, the most elegant of which uses the Gegenbauer expansion.

Read more about this research using this link.

4. The Spike, the Sparse, and the Sink

This research paper studies two phenomena that usually co-occur and often involve the same tokens in Transformer-based LLMs:

Massive activations: where some tokens produce extremely large activations (“Spikes”)
Attention sinks: where some tokens consistently attract attention disproportionate to their semantic importance (“Sinks”)

Experiments show that:

The co-occurrence of the two is an architectural artifact of modern Transformer design.
Massive activations work globally to create nearly constant hidden representations that persist across many layers, effectively acting like implicit parameters of the model.
Attention sinks work locally by influencing attention outputs across heads, pushing individual heads to focus more on short-range dependencies.
Pre-norm is the main architectural choice that allows these two phenomena to appear together, and removing pre-norm causes them to decouple.

Read more about this research using this link.

5. Speculative Speculative Decoding

Speculative decoding has become a standard way to accelerate autoregressive inference in LLMs.

It uses a fast draft model to predict upcoming tokens from a slower, larger target model, and then verifies them in parallel with a single forward pass of the target model.

However, the method is still partly sequential because the draft must wait for verification before making the next speculation.

This research paper introduces Speculative speculative decoding (SSD), which parallelizes the speculation and verification operations.

While verification is ongoing, the draft model predicts likely verification outcomes and prepares pre-emptive speculations for them. If the actual verification outcome is then in the predicted set, a speculation is returned immediately, removing drafting delays entirely.

This method is further improved, leading to Saguaro, an optimized SSD algorithm that is up to 5x faster than autoregressive decoding on open-source inference engines.

Read more about this research using this link.

6. Olmo Hybrid

Hybrid LLMs use an architecture that combines transformer attention and linear RNN layers.

Despite previous releases such as Nemotron-H, Qwen3-Next, Kimi Linear, and Qwen 3.5, it was unclear if the purported benefits of hybrid architectures justify the cost of scaling them up.

This debate has been settled by Ai2 researchers, who released the fully open-source Olmo Hybrid, a 7B-parameter model largely comparable to Olmo 3 7B, but with its sliding-window layers replaced by Gated DeltaNet layers.

Olmo Hybrid is trained using 49% fewer training tokens and outperforms Olmo 3 across standard pre-training and mid-training evaluations.

This is empirical proof that hybrid architectures not only reduce memory usage during inference but also obtain more expressive models that scale better during pretraining.

Architecture comparison for Olmo 3 and Olmo Hybrid. Olmo Hybrid replaces the sliding-window attention layers from Olmo 3 with Gated DeltaNet.

Read more about this research using this link.

7. Reasoning Models Struggle to Control their Chains of Thought

Chain-of-thought (CoT) monitoring is a promising method for understanding the intents and finding misbehaviors in modern reasoning LLMs.

This research paper examines whether CoT monitoring is a reliable safety tool or whether reasoning models can control what they reveal in their chain-of-thought (CoT), a phenomenon known as CoT controllability.

For this, the researchers introduce an evaluation suite called CoT-Control, which includes tasks that require reasoning models to solve problems while adhering to CoT instructions. (For example, reasoning about a genetics question without using the word “chromosome”.)

Experiments show that:

Reasoning models have significantly lower CoT controllability than output controllability. For example, Claude Sonnet 4.5 can control its CoT only 2.7% of the time, but 61.9% of the time when controlling its final output.
CoT controllability is higher for larger models and decreases with more RL training, test-time compute, and increased problem difficulty.
CoT controllability failures extend even to situations in which models are given incentives (as opposed to direct requests) to evade CoT monitors
Models exhibit slightly higher controllability when they are told they are being monitored.
Eliciting controllability by adversarially optimizing prompts does not meaningfully increase controllability.

The overall results indicate that CoT controllability is currently unlikely to be a failure mode of CoT monitorability and must be tracked by frontier labs in future models.

Read more about this research using this link.

8. Phi-4-reasoning-vision-15B

This research paper from Microsoft introduces Phi-4-reasoning-vision-15B, a 15B-parameter multimodal reasoning model that can process both text and images.

The model combines the Phi-4-Reasoning backbone with a SigLIP-2 vision encoder in a mid-fusion architecture that converts images into visual tokens for reasoning.

It performs efficient multimodal reasoning across tasks such as math, science, OCR, document understanding, and GUI interaction. This is made possible by the use of high-resolution visual encoder and by training the model on a mix of reasoning and non-reasoning data with explicit mode tokens. This setup allows the model to provide fast, direct answers for simple tasks and chain-of-thought reasoning for more complex ones.

The paper shows that the greatest performance improvements come from systematic data filtering, error correction, and the augmentation with synthetic data. This tells us that data quality is the main factor behind how well the model performs.

Read more about this research using this link.

9. Helios

This research paper presents Helios, the first 14B autoregressive diffusion model that can create high-quality minute-scale videos in real time at about 19.5 FPS on a single NVIDIA H100 GPU.

Helios also supports text-to-video, image-to-video, and video-to-video generation within a single framework, and extensive experiments show that it consistently outperforms prior methods on both short and long-video generation.

Read more about this research using this link.

10. Symbol-Equivariant Recurrent Reasoning Models

This research paper introduces Symbol-Equivariant Recurrent Reasoning Models (SE-RRMs), a new addition to the family of Recurrent Reasoning Models (RRMs).

SE-RRMs extend Recurrent Reasoning Models by explicitly encoding symbol-permutation symmetry in the architecture using symbol-equivariant layers.

This change ensures that the model provides the same solution even when symbols, such as numbers or colors, are permuted. This reduces the need for extensive data augmentation and improves generalization.

Experiments show that SE-RRMs outperform previous RRM models on Sudoku and can adapt to different Sudoku grid sizes, which existing RRMs cannot.

SE-RRMs also achieve competitive performance on ARC-AGI-1 and ARC-AGI-2 with substantially less data augmentation and only 2M parameters.

Read more about this research using this link.

Read more about the previous RRMs using the following links:

This article is entirely free to read. If you loved reading it, restack and share it with others to earn referral rewards. ❤️

Don’t forget to become a paid subscriber at a 25% discount to access all valuable posts in the newsletter.

Your 25% off

Learn to train a deep learning model on multiple GPUs with Distributed PyTorch

Dr. Ashish Bamania — Sun, 08 Mar 2026 11:33:47 GMT

👋🏻 Hey there!

This is the second part of a two-part lesson where we get our hands dirty by writing PyTorch code to train a model across multiple GPUs using the distributed machine learning algorithm Distributed Data Parallel (DDP).

If you missed the first part, where we learned about DDP visually, here is your link to get started.

DDP sounds straightforward conceptually, but the intricate dependencies between GPU computation and communication make it tricky to implement and optimize.

Fortunately, we do not have to rewrite it from scratch and handle all existing edge cases. PyTorch already implements all of these in its PyTorch Distributed library in the DistributedDataParallel class.

To make this lesson easy to follow, I will write all the code in a Jupyter notebook on Kaggle.

Kaggle provides free access to two NVIDIA T4 GPUs for 30 hours each week. We will train a CNN on these two GPUs with DDP using the CIFAR-10 dataset.

The CIFAR-10 dataset contains 60,000 32x32 colour images in 10 classes, with 6000 images per class. Some examples from the dataset are shown in this image.

Let’s begin!

Blessings from Jensen Huang (Source)

Import packages

We start with importing the necessary packages as follows.

# PyTorch core
import torch                              
import torch.nn as nn                    
import torch.optim as optim  
from torch.utils.data import DataLoader                 

# PyTorch Distributed
import torch.distributed as dist          
import torch.multiprocessing as mp        
from torch.nn.parallel import DistributedDataParallel as DDP  
from torch.utils.data.distributed import DistributedSampler  

# TorchVision
import torchvision                        
import torchvision.transforms as transforms 

# Standard library utils
import os

Define the CNN classifier

Our CNN takes an image as input and classifies it into one of the 10 classes.

We define it exactly the same way we would when training on a single machine, and we aren’t doing anything fancy in this step.

class ImageClassifier(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        
        self.cnn = nn.Sequential(
            # Feature extraction
            # Block 1
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=64),
            nn.ReLU(),

            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=64),
            nn.ReLU(),

            nn.MaxPool2d(kernel_size=2),

            # Block 2
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=128),
            nn.ReLU(),

            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=128),
            nn.ReLU(),

            nn.MaxPool2d(kernel_size=2),

            # Block 3
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=256),
            nn.ReLU(),

            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=256),
            nn.ReLU(),

            nn.MaxPool2d(kernel_size=2),
            
            nn.Flatten(),
            nn.Linear(in_features=256 * 4 * 4, out_features=512),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            
            # Classification
            nn.Linear(in_features=512, out_features=num_classes),
        )

    def forward(self, x):
        return self.cnn(x)

Write helper functions for distributed training

Now comes the distributed part.

A Simple 3-Step Guide to Move Your Memories From ChatGPT to Claude

Dr. Ashish Bamania — Wed, 04 Mar 2026 16:21:36 GMT

Millions decided to quit ChatGPT this week, given OpenAI’s new deal with the US Department of War.

Their first choice was Claude, given that Anthropic did not participate in this deal and refused to use its AI models for mass surveillance and developing fully autonomous weapons.

If you’re making this move as well, don’t forget to move your memories into Claude. Here’s a simple 3-step guide to do this.

Note that:

Claude’s memory is also available on the free plan.
This guide isn’t ChatGPT-specific and will let you move memories from any LLM provider to Claude.

Step 1

Paste the following prompt when you’re chatting with the LLM that you want to move away from.

I’m moving to another service and need to export my data. List every memory you have stored about me, as well as any context you’ve learned about me from past conversations. Output everything in a single code block so I can easily copy it. Format each entry as: [date saved, if available] - memory content. Make sure to cover all of the following — preserve my words verbatim where possible: Instructions I’ve given you about how to respond (tone, format, style, ‘always do X’, ‘never do Y’). Personal details: name, location, job, family, interests. Projects, goals, and recurring topics. Tools, languages, and frameworks I use. Preferences and corrections I’ve made to your behavior. Any other stored context not covered above. Do not summarize, group, or omit any entries. After the code block, confirm whether that is the complete set or if any remain.

Step 2

Copy all the results that you obtain in the chat response.

Step 3

Go to your Claude Memory settings, paste the copied result in the text box, and click the “Add to memory” button.

That’s it!

That’s all you need to do to transfer memories, so Claude knows about you and gives responses best suited to you.

If you have any other questions, let me know in the comments below.

Bye now! 👋🏻

🤗 I want your help to make your life better

Dr. Ashish Bamania — Wed, 04 Mar 2026 13:38:09 GMT

👋🏻 Hey friend!

I’ll keep this short.

I want to make “Into AI” as useful as possible for your career, and the best way to do that is to ask you directly.

What topics are you most interested in right now?
What’s the biggest challenge you’re facing in your career?
What kind of content would actually help you move forward?

Just hit reply to this email and let me know. 📧

There’s no wrong answer.

Even a one-liner is totally fine.

I read every single reply, and your answers will directly shape what comes next in this newsletter.

🎁 As a thank you, I’ll send you a personalized 40% discount coupon on the annual subscription just for sharing your thoughts.

Talk soon!

Ashish

🗓️ This Week In AI Research (22-28 February 26)

Dr. Ashish Bamania — Tue, 03 Mar 2026 10:11:35 GMT

🎁 Here’s your 25% off the annual membership. Unlock every lesson in the newsletter and accelerate your AI engineering career!

Claim your discount now!

1. ChatGPT Health performance in a structured test of triage recommendations

This research paper, published in Nature Medicine, evaluated ChatGPT Health, a popularly used AI tool for triaging patients.

The study found that it failed to appropriately triage over half (52%) of serious emergencies like diabetic ketoacidosis or impending respiratory failure and suggested a follow-up in 24 to 48 hours instead of seeking immediate emergency care.

Additionally, it mismanaged patients with suicidal ideations, sometimes triggering alerts in low-risk cases and failing to respond when individuals mentioned specific self-harm methods (which actually made them high-risk clinically).

Its recommendations also shifted toward less urgent care when relatives and friends downplayed symptom severity.

Read more about this research using this link.

2. Mercury 2

Inception announced Mercury 2, the world's fastest reasoning language model, which uses diffusion-based generation rather than the traditional autoregressive generation.

This allows it to generate and refine many tokens in parallel, resulting in much lower latency in production workflows.

Mercury 2 generates over 1,000 tokens per second on NVIDIA Blackwell GPUs, making it more than five times faster than typical speed-optimized models while maintaining competitive generation quality.

It comes with a 128K context window, tunable reasoning effort, native tool use, and schema-aligned JSON output, and is OpenAI API compatible, so it can be dropped into existing stacks without requiring any rewrites.

Read more about this release using this link.

3. DualPath: Breaking the Storage Bandwidth Bottleneck in Agentic LLM Inference

Multi-turn, agentic LLM workloads are limited by storage I/O during the loading of large key-value (KV) caches rather than computation.

This issue is particularly common in popular architectures, where loading large KV caches from external storage creates an imbalance, with prefill engines consuming their storage network bandwidth while decode engines are not fully utilized. This restricts the system's overall throughput.

To solve this problem, the authors introduce DualPath, a new inference system that creates a second path for loading KV caches.

Alongside the usual storage-to-prefill path, it creates a new storage-to-decode path that loads KV caches into decoding engines. These caches are then transferred via RDMA to prefill engines.

This optimized data path is combined with a global scheduler that dynamically balances load across prefill and decode engines.

Experiments on three models with realistic agentic workloads show that DualPath improves offline inference by up to roughly 1.87× and online serving throughput by about 1.96×, all while meeting the SLO.

Read more about this research using this link.

4. Agents of Chaos

This research paper presents a read-teaming study of autonomous LLM-based agents deployed in a live lab environment with persistent memory, email accounts, Discord, file systems, and shell execution.

The goal was to observe how these agents interact with humans and other agents under benign and adversarial conditions. Over two weeks, twenty AI researchers tested the systems and recorded eleven significant failure cases.

These included:

Unauthorized compliance with non-owners
Disclosure of sensitive information
Execution of destructive system-level actions
Denial-of-service conditions
Uncontrolled resource consumption
Identity spoofing vulnerabilities
Cross-agent propagation of unsafe practices
Partial system takeover
Fake reporting of task completion

These findings reveal major issues with agentic systems when deployed in the real world.

Read more about this research using this link.

5. dLLM: Simple Diffusion Language Modeling

This research paper introduces dLLM, an open-source framework that unifies the core components of diffusion language modeling (training, inference, and evaluation) and makes them easy to customize.

With dLLM, users can reproduce, finetune, deploy, and evaluate open-source large DLMs such as LLaDA and Dream through a standardized pipeline.

The framework also provides minimal, reproducible recipes for building small DLMs from scratch with accessible compute, including converting any BERT-style encoder or autoregressive LM into a DLM.

The researchers release the checkpoints of these small DLMs to make DLMs more accessible and accelerate future research.

Read more about this research using this link.

6. Generative Recommendation for Large-Scale Advertising

This research paper introduces GR4AD (Generative Recommendation for ADdvertising), a generative recommendation system focused on production for large-scale advertising, where real-time performance and business value are important.

GR4AD introduces:

UA-SID for tokenization to capture complicated business information, and
LazyAR, a lazy autoregressive decoder that lowers inference costs for short, multi-candidate outputs

To better align optimization with business metrics, it uses Value-Aware Supervised Learning (VSL) and a ranking-aware reinforcement learning algorithm called RSPO (Ranking-Guided Softmax Preference Optimization).

For online inference, a method called Dynamic Beam Serving is used that adjusts beam width according to generation workload to balance quality and compute power.

In large-scale online A/B tests, GR4AD achieved up to a 4.2% increase in ad revenue compared to a baseline recommendation stack.

It has also been fully deployed in Kuaishou’s advertising system, serving over 400 million users and achieving high-throughput, real-time serving.

Read more about this research using this link.

7. Aletheia tackles FirstProof autonomously

This research paper from Google DeepMind reports the performance of Aletheia, a mathematics research agent powered by Gemini 3 Deep Think, on the FirstProof challenge, a collection of 10 research-level math problems with publicly undisclosed solutions used to assess AI capability.

Aletheia solved 6 out of 10 problems within the allowed time, with most expert evaluators considering those solutions correct. Problem 8 had mixed expert agreement.

Read more about this research using this link.

8. LeRobot: An Open-Source Library for End-to-End Robot Learning

This research paper introduces LeRobot, an open-source library that brings together and simplifies the entire robot-learning stack.

It covers everything from low-level middleware for controlling robots to large-scale data collection, storage, streaming, and learning algorithms, within a single framework.

Instead of relying on hand-crafted techniques, it introduces scalable methods that lower the barriers for robotics researchers and practitioners worldwide.

Read more about this research using this link.

9. Training Agents to Self-Report Misbehavior

LLM agents may pursue misaligned goals while hiding their actions, despite being Alignment-trained.

To counter this, this research paper introduces Self-incrimination training, which trains agents to produce a visible signal when they act deceptively.

The researchers train agents based on GPT-4.1 and Gemini-2.0 to use a report_scheming() tool whenever they behave inappropriately and measure their ability to cause harm undetected in out-of-distribution environments.

Self-incrimination significantly reduces undetected attacks and outperforms similar monitoring and alignment methods, while preserving the instruction hierarchy and causing little loss in capability.

Its effectiveness stays steady across tasks, continues to work under adversarial prompt optimization, and applies to situations where agents pursue misaligned goals themselves rather than being instructed to misbehave.

Read more about this research using this link.

10. On Data Engineering for Scaling LLM Terminal Capabilities

This research paper from NVIDIA is a systematic study of data engineering practices to build terminal agents. It introduces:

Terminal-Task-Gen: a lightweight synthetic task generation pipeline that supports seed-based and skill-based task construction
Terminal-Corpus: a large-scale open-source dataset for terminal tasks

Using this dataset, the researchers train Nemotron-Terminal, a family of models initialized from Qwen3 that achieve substantial gains on Terminal-Bench 2.0, matching the performance of significantly larger models.

Read more about this research using this link.

11. How to Train Your Deep Research Agent?

This research paper explores how to best train Deep Research agents.

By studying three RL training choices (Prompt templates, Reward functions, and Policy optimization), the researchers reveal three main findings:

Fast Thinking prompts work better than Slow Thinking
F1-based rewards need action penalties to avoid collapse and can outperform exact match (EM)
REINFORCE outperforms PPO while requiring fewer search actions, whereas GRPO shows the poorest stability among policy optimization methods

Building on these findings, the researchers introduce Search-R1++, a strong baseline that improves the performance of Search-R1 from 0.403 to 0.442 (Qwen2.5-7B) and 0.289 to 0.331 (Qwen2.5-3B).

Read more about this research using this link.

If you loved reading it, restack and share it with others to earn referral rewards. ❤️

Become a paid subscriber today to get access to all posts on this newsletter.

You can also check out my books on Gumroad and connect with me on LinkedIn to stay in touch.

11 Infographics on the Dark Side of AI Kept Hidden From You

Dr. Ashish Bamania — Sun, 01 Mar 2026 13:07:55 GMT

👋🏻 Hey friends!

Artificial Intelligence is one of the greatest technologies of our lifetime. But there’s a dark side to it which you might not be aware of.

In this article, I present 11 infographics that tell the flaws and risks associated with current AI systems.

Get ready to have your mind blown! 🤯

🤝 Before we begin, I want to introduce you to Emma McIlwaine, the brilliant graphic designer and illustrator behind this newsletter’s infographics.

Check out her designs on her website, follow her on LinkedIn, and reach out to work with her.

1. ChatGPT Health would cause serious harm if deployed in hospitals

ChatGPT Health today answers health queries for millions of people around the world, but is it ready to do so?

Maybe not. A recent study published in Nature Medicine showed that:

It under-traiged 52% of medical emergencies, such as Diabetic ketoacidosis and impending Respiratory failure, to a 24–48-hour evaluation rather than the emergency department.
When a patient’s symptoms were downplayed by friends or family, the triage recommendations shifted towards less urgent care.
When people expressed suicidal thoughts with a specific method (which means a higher risk clinically), crisis interventional messages to get urgent help were less likely to be shown.

2. AI is drinking the world dry

According to the World Bank, 4 billion people live in water-scarce areas today.

By 2030, global demand for water will exceed sustainable supply by as much as 40%, and 1.6 billion people will not have access to safe drinking water.

At present, the AI economy requires 23 cubic km of water a year. By 2050, this is predicted to more than double to more than 54 cubic kilometers.

Unfortunately, this means that our world will run out and will need to find an extra 31 cubic kilometers of water per year to run the AI economy. That’s enough to supply every human being on Earth with an extra 3,820 liters of freshwater a year.

According to another study, training GPT-3 could have consumed a total of 5.4 million liters of water, with 700,000 liters of clean freshwater evaporated directly in the training data centers.

Alongside this, GPT-3 needs to “drink” a 500ml bottle of water for roughly 10-50 medium-length responses, depending on when and where it is deployed.

Information about the water consumption of our current AI systems is largely kept a secret.

3. AI has a race problem

A study from Lehigh University found that LLMs consistently recommended denying more loans and charging higher interest rates to Black applicants than to otherwise identical white applicants.

The study also found that black applicants would need credit scores ~120 points higher than white applicants to receive the same approval rate, and ~30 points higher to receive the same interest rate.

Another study from the University of Washington revealed that AI resume screeners preferred white-associated names in 85.1% of cases, and Black male candidates were preferred 0% of the time.

Read it again: It’s 0%.

4. Copyright? What’s that?

Anthropic recently released a detailed report on how three Chinese AI labs used ~24,000 fraudulent accounts to extract their state-of-the-art AI model, Claude’s capabilities to improve their own models.

But has Anthopic always been fair?

A recent study revealed that current LLMs can narrate almost the entire text of ‘Harry Potter and the Sorcerer’s Stone’ word for word, with Claude 3.7 Sonnet reaching the highest recall of 95.8%, followed by Gemini 2.5 Pro (76.8%) and Grok 3 (70.3%).

5. Yes, AI is really coming for your job

Former CEO of Twitter/ X and current CEO of Block, Jack Dorsey, fired 4,000 employees this week, saying AI can work faster and better.

“We're already seeing that the intelligence tools we’re creating and using, paired with smaller and flatter teams, are enabling a new way of working which fundamentally changes what it means to build and run a company. and that's accelerating rapidly. — Jack”

@blocks smaller today. here's my note to the company.\n\n####\n\ntoday we're making one of the hardest decisions in the history of our company: we're reducing our organization by nearly half, from over 10,000 people to just under 6,000. that means over 4,000 of you are","username":"jack","name":"jack","profile_image_url":"https://pbs.substack.com/profile_images/1661201415899951105/azNjKOSH_normal.jpg","date":"2026-02-26T21:12:36.000Z","photos":[],"quoted_tweet":{},"reply_count":8535,"retweet_count":6479,"like_count":49292,"impression_count":57682210,"expanded_url":null,"video_url":null,"belowTheFold":true}" data-component-name="Twitter2ToDOM">

But this isn’t a one-off event.

The IMF says that almost 40% of global employment is exposed to AI, rising to 60% in advanced economies.

Moreover, reports from Challenger, Gray & Christmas clearly show that AI caused nearly 55,000 job cuts in the US in 2025 alone.

6. Your passport is in AI training data

Data privacy in the AI age is largely a myth.

A study that examined 0.1% of a popular large-scale web-scraped dataset used for training AI systems called CommonPool, found millions of images of passports, credit cards, and birth certificates, with an estimated 102 million unblurred faces.

Even examining every sample in the 0.1% was extremely difficult for the researchers, given the size of this dataset. Imagine what the remaining 99.9% could contain?

7. AI is choosing who lives and dies in war

Anthropic has been in debates with the U.S. Department of War over using its AI models to develop systems for mass domestic surveillance and fully autonomous weapons.

That might be great news for America, but many other countries have gone beyond such discussions.

(Update: I’m sorry to tell you that OpenAI has now accepted the above proposal.)

Israel’s defense forces reported using AI systems known as "Lavender" to identify targets in Gaza, and "The Gospel" to help select and prioritize physical strike targets.

According to the +972 Magazine, Israel’s army almost completely relied on "Lavender," which flagged 37,000 Palestinians as suspected militants and their homes for possible air strikes.

Despite knowing that the system makes errors in approximately 10% of cases, there was no requirement to check why the AI system made its choices.

Once the target was identified by “Lavender”, human personnel would take just 20 seconds to authorize a bombing to ensure the marked target was male.

8. Researchers are using AI to write fake science

AI-written papers are becoming increasingly common in academia. To add to this, many are even using LLMs to peer-review journal submissions.

A study published in Nature showed that more than 10,000 research papers were retracted in 2023, setting a new record. AI-generated fraud was found to be one of the increasingly common reasons for these retractions.

What happens when results from these fake research papers find their way to products that directly affect your life?

9. Yes, AI might now send you to jail

"AI won't take over for people. But people who use AI will take the place of those who don't.” — Sam Altman

Many have held on to this quote quite passionately despite knowing what AI systems can and cannot do. This includes lawyers as well.

There have been more than 900 cases worldwide where AI-generated hallucinated content appeared in court filings.

According to another study by Stanford HAI, general-purpose AI chatbots hallucinate between 58% and 82% of the time in responses to legal queries.

A single hallucination like this is enough to cause a wrong legal judgment and send someone to jail.

God bless us all.

10. The world lives in the dark as AI electrifies

Data from the International Energy Agency (IEA) shows that 730 million people worldwide still lacked access to electricity in 2024.

Global data centers, on the other hand, consumed 415 TWh of electricity in 2024 and are projected to reach 945 TWh by 2030. This is more than Japan's total electricity consumption today!

11. Deepfakes are here to steal your identity

Deepfakes sounded like a thing of the future a few years ago, but such a dystopian future is already here.

According to a report from the European Parliament Research Service (EPRS), Deepfakes grew from ~500,000 in 2023 to a projected 8 million in 2025.

98% of them were non-consensual pornography and mostly targeted women.

If you loved reading this article and found it valuable, restack to share it with others and earn some referral rewards. ❤️

Don’t forget to connect with Emma on LinkedIn and reach out to work with her.

If you want to get even more value from this publication, become a paid subscriber to get access to posts like the following and more:

Bye now! 👋🏻

Learn to train deep learning models on multiple GPUs

Dr. Ashish Bamania — Thu, 26 Feb 2026 13:06:02 GMT

We have long ago moved away from training production-grade deep learning models on a single machine.

Today, large deep learning models such as LLMs are usually trained for months on distributed clusters with multiple GPUs. This is because a single machine usually lacks the memory and computing power needed to process massive training datasets and train a model.

For example, Llama 3.1 models were pre-trained on around 15 trillion data tokens and required a total of 39.3M GPU hours of computation on NVIDIA H100-80GB GPUs. That’s a big number!

LLM training in the times of the Flintstones (Source)

In this two-part lesson, we will learn about an algorithm called Distributed Data Parallel (DDP), and then get our hands dirty by writing some PyTorch code to train a model across multiple GPUs using it.

Let’s begin!

Btw, I’m running a 50% off sale on the annual subscription to ‘Into AI’, valid ONLY for the next 2 days.

Grab your discount now!

Claim your discount NOW and get access to articles like these, and so much more:

But first, what is Distributed Data Parallel (DDP)?

Distributed data parallel (DDP) is a popular algorithm for training models across multiple GPUs.

The algorithm uses Data parallelism, which means splitting the training dataset across multiple GPUs so each GPU processes a different portion of it in parallel during model training.

Here are the steps that take place in DDP:

The model is first replicated on each GPU along with a copy of the optimizer.
The training dataset is split across GPUs using DistributedSampler, which ensures that each GPU trains on a unique data subset/ batch per epoch.
Each GPU runs a forward pass on its subset of the training dataset to calculate loss, and stores the activations
Each GPU runs a backward pass using the activations and loss, and computes gradients
All GPUs communicate and share their calculated gradients with each other
The gradients are averaged across the GPUs using an algorithm called AllReduce
Averaged gradients are communicated back to each GPU
Each GPU updates its copy of model weights using the averaged gradients. As each GPU starts with the same model weights and updates them using the same averaged gradients, all the resulting model copies remain in sync with other GPUs.
The next batch of training data is loaded, and the cycle repeats from step 3 till the training epoch ends.
In case of training for multiple epochs, the entire training dataset is reshuffled using DistributedSampler and is iterated over again.
Once all training epochs end, the trained model is returned.

DDP makes the training process even more efficient by running steps 4 and 5 at the same time.

It doesn't wait for gradients from all layers to be computed before sharing them. Instead, it sends each layer's gradients to the other GPUs as soon as they're ready, so that communication occurs in parallel with computation.

Understand DDP Visually

Here is an example of a 3-GPU distributed training setup using DDP.

In the illustrations shown below:

Each GPU starts with its own copy of the model (and optimizer).
Then, each GPU is given different samples/ batch from the training dataset
Each GPU runs a forward and backward pass, computing loss and gradients

Following this, gradients from all GPUs are aggregated and averaged using the AllReduce operation.

After this step, the averaged gradients are broadcast back to all the GPUs.
Each GPU then updates its copy of model weights using them. This keeps all copies of the model between different GPUs synchronized.

Once the last training epoch ends, the trained model is returned.

But DDP Isn’t Perfect

DDP speeds up and efficiently scales model training, but it is not a one-stop solution to distributed model training.

DDP involves each GPU storing at least:

Model parameters
Optimizer and its state
Model activations after forward pass
Gradients after backward pass

Consider the NVIDIA H100 (80GB) GPU, which is widely used for model training.

If we train an 8B-parameter Llama 3 model on it using the AdamW optimizer, with full-precision training (which means that all numbers used are FP32 or require 4 bytes):

The model parameters will require: 8B × 4 bytes = 32 GB
Gradients will require: 8B × 4 bytes = 32 GB
The AdamW optimizer state (first and second moments) will require: 8B × 4 bytes x 2 = 64 GB
Alongside these, activations will require 10-20 GB, with the use of activation checkpointing
Other unaccounted computation/ communication overheads might require another 2-3 GB

The total requirements of 151 GB far exceed the GPU's VRAM capacity of 80 GB. This means that training such multi-billion-parameter models isn’t possible with DDP.

This is when we move towards using Model Parallelism techniques, which involve splitting model parameters across multiple GPUs. This is a topic of future lessons in the newsletter.

In the next part of this lesson on DDP, we will learn how to train a deep learning model using it on free-to-use cloud GPUs.

If you loved reading it and found it valuable, restack to share it with others. ❤️

Don’t forget to grab your 50% off sale on the annual subscription to ‘Into AI’, valid ONLY for the next 2 days.

Grab your discount now!

🗓️ This Week In AI Research (15-21 February 26)

Dr. Ashish Bamania — Mon, 23 Feb 2026 14:47:18 GMT

🎁 Before we start, I want to offer you a 35% DISCOUNT on the annual subscription to ‘Into AI’ as part of the Fab Feb offer. It’s only valid till the end of February, so don’t wait up!

1. Gemini 3.1 Pro

Google released Gemini 3.1 Pro, a new model in the Gemini 3 series with advanced reasoning capabilities for complex problem-solving.

3.1 Pro achieved a verified score of 77.1% on ARC-AGI-2. This is more than double the reasoning performance of its predecessor, 3 Pro.

Read more about this release using this link.

2. MemoryArena

Current evaluation methods of agents test either one of the following in isolation:

Memorization: Recall of past conversations, but not capturing how memory is used to guide future decisions
Action: Agent acting in single-session tasks without the need for long-term memory

In reality, memorization and action are tightly coupled. MemoryArena is a new benchmark that assesses this setting by testing agents' memory in multi-session tasks that require remembering and using information from earlier sessions to succeed in later ones.

MemoryArena covers four domains:

Bundled Web Shopping
Group Travel Planning
Progressive Web Search, and
Sequential formal reasoning in math/ physics

Results show that while many agents might perform flawlessly on existing memory benchmarks, such as LoCoMo, they do not perform well on MemoryArena.

Alongside this, external memory and RAG systems do not consistently improve the performance of these agentic systems. This reveals a major issue that current benchmarks do not adequately test agents in real-world settings.

Read more about this research paper using this link.

3. Fast KV Compaction via Attention Matching

Scaling LLMs to long contexts is bottlenecked by the size of the KV cache. To address this, in deployed models, long contexts are usually compacted using summarization or token dropping, which degrades performance.

This research paper addresses this problem by introducing a method for fast context compaction in latent space using Attention Matching.

This constructs compact keys and values to reproduce attention outputs and preserve attention mass at a per-KV-head level.

The approach breaks down into efficient subproblems with many closed-form solutions, which enables KV compaction ratios of up to 50x in seconds with minimal loss in quality. These results are orders of magnitude faster than previous latent-space methods.

Read more about this research paper using this link.

4. Towards a Science of AI Agent Reliability

The difference between AI agents' performance on benchmarks and in the real world is quite striking. This is because current benchmarks test agent behavior using a single success metric (accuracy), which hides many of their critical operational flaws.

This research paper addresses this by proposing 12 specific metrics that test agent reliability along four key dimensions:

Consistency
Robustness
Predictability
Safety

An evaluation of 14 agentic models across two complementary benchmarks shows that recent capability gains yield only small improvements in reliability, with consistency and predictability lagging the most.

Read more about this research using this link.

Grab your copy today at a 20% discount using this link.

5. Experiential Reinforcement Learning

This research paper introduces Experiential Reinforcement Learning (ERL), a new training approach for LLMs that includes an explicit experience–reflection–consolidation loop into the RL process.

For a given task, ERL has a model generate its first attempt. The model then receives environmental feedback and generates a self-reflection conditioned on this attempt, which guides it to produce a more refined second attempt.

Both attempts and reflections are optimized with RL, while successful second attempts are internalized via self-distillation, so the model learns to reproduce improved behavior directly from the original input without self-reflection.

Across sparse-reward control environments and agentic reasoning benchmarks, ERL improves learning efficiency and final performance over strong RL baselines.

The gains are of up to +81% in complex multi-step environments and up to +11% in tool-using reasoning tasks.

Read more about this research using this link.

6. Calibrate-Then-Act

Solving complex problems requires LLM-based agents to interact with an environment to gather information.

This research paper introduces a framework that enables agents to balance this exploration cost and uncertainty when interacting with an environment, thereby performing more optimal environment exploration.

Multiple tasks are formalized as sequential decision-making problems under uncertainty. Each problem has a latent environment state that can be reasoned about via a prior. The framework, called Calibrate-Then-Act (CTA), then feeds the LLM this additional context to enable it to act more optimally.

Results on information-seeking QA and on a simplified coding task show that making cost-benefit tradeoffs explicit with CTA can help agents discover more optimal decision-making strategies.

Read more about this research using this link.

7. Evaluating Collective Behaviour of Hundreds of LLM Agents

This research paper presents a scalable evaluation framework to study how large groups of LLM-based autonomous agents behave in social dilemmas.

With this framework, LLMs generate strategies encoded as algorithms, enabling their inspection prior to deployment.

Experiments show that newer models often lead to worse social outcomes than older ones when agents focus on personal gain instead of group benefits.

By using cultural evolution simulations to model user selection of agents, researchers find a significant risk of poor social outcomes, especially as the relative benefit of cooperation decreases, and population size grows.

Read more about this research using this link.

8. ZUNA

This research paper from Zyphra introduces ZUNA, an open-sourced foundation model with 380 million parameters for EEG (electroencephalogram) brain-computer interface (BCI) data.

The model is designed to denoise, reconstruct, and upsample EEG signals from different electrode layouts.

It is built as a diffusion autoencoder that can take a subset of EEG channels and reliably recover missing or corrupted ones. This improves spatial resolution and signal quality for research and practical uses.

ZUNA can work across various datasets and hardware, and serves as a foundation for non-invasive “thought-to-text” systems and for better understanding neural signals.

Read more about this research using this link.

9. KLong

This research paper introduces KLong, an open-source LLM agent trained to solve extremely long-horizon tasks. These are tasks that would exceed the context window without context management and involve long-running experiments such as those from MLE-bench and PaperBench.

Experiments show that KLong greatly enhances performance on long-horizon benchmarks compared to baselines.

Notably, KLong (106B) beats Kimi K2 Thinking (1T) by 11.28% on PaperBench. This performance improvement generalizes to other coding benchmarks, such as SWE-bench-Verified and MLE-bench.

Read more about this research using this link.

10. Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

This research paper presents NAMO and NAMO-D, two new optimization algorithms that combine the benefits of orthogonal updates (as seen in Muon optimizer) with the noise adjustment of Adam-type methods.

NAMO preserves momentum orthogonality and introduces a single adaptive learning rate, which enhances Muon without any significant extra cost.

Its diagonal variant, NAMO-D, further improves performance by allowing neuron-wise noise adaptation through a clamped diagonal scaling.

The research paper includes theoretical convergence guarantees in both deterministic and stochastic environments, showing that both optimizers adapt well to gradient noise.

Empirical results from GPT-2 pretraining also show that NAMO, particularly NAMO-D, surpasses both AdamW and Muon baselines.

Read more about this research paper using this link.

This article is entirely free to read. If you loved reading it, restack and share it with others to earn referral rewards. ❤️

Don’t forget to become a paid subscriber with a 35% discount on the annual subscription. Offer ends this month!

Grab your offer NOW!

You can also check out my books on Gumroad and connect with me on LinkedIn to stay in touch.

Build a Personal Health Companion with Mem0 and CrewAI

Dr. Ashish Bamania — Fri, 20 Feb 2026 11:06:31 GMT

Hey there!

In the last few lessons, we learned about AI memory.

In this lesson, we will apply these concepts to build a multi-agent system that can monitor personal health issues/comorbidities and act as a long-term health companion.

The tools that we will use are:

CrewAI: Framework for building multi-agent systems
Mem0: Persistent memory layer for LLMs

If you’re completely new to building agents with CrewAI, the following lesson will help you get started.

Let’s begin!

Overview of the Multi-Agent System

We intend to build a health companion that manages chronic health conditions/ comorbidities using three specialized AI agents that work together sequentially.

These agents are:

🗓️ This Week In AI Research (8-14 February 26)

Dr. Ashish Bamania — Mon, 16 Feb 2026 12:30:12 GMT

Before we start, I want to introduce you to my book, LLMs In 100 Images, which is a collection of 100 easy-to-follow visuals that explain the most important concepts you need to master to understand LLMs today.

1. GPT‑5.2 derives a new result in theoretical physics

OpenAI’s GPT‑5.2 helped develop a new theoretical physics result by proposing a simpler formula for a previously overlooked gluon scattering amplitude.

This configuration, once thought to be zero under standard assumptions, actually has a non-zero amplitude in a special kinematic regime.

The formula first proposed by GPT‑5.2 was later verified using another internal OpenAI model and by human researchers.

Read more about this achievement using this link.

Read the physics research paper using this link.

2. Towards Autonomous Mathematics Research

This research from Google DeepMind introduces Aletheia, a math research agent designed to move beyond Olympiad-style problem-solving into professional mathematics research, where the big challenges include navigating vast literature and constructing long-term proofs.

Aletheia can iteratively generate, verify, and revise complete solutions in natural language and uses:

An advanced version of Gemini Deep Think for hard reasoning
A new inference-time scaling law that extends beyond Olympiad-level problems
Intensive tool use to handle real mathematical research complexities

With these, Athelia has achieved three significant research milestones:

A fully AI-generated paper on calculating structure constants in arithmetic geometry, called ‘Eigenweights’.
A human-AI collaborative paper that proves bounds on systems of interacting particles called ‘independent sets’.
A semi-autonomous evaluation of 700 open problems in Bloom’s Erdős Conjectures database, which led to autonomous solutions for four open questions.

Read more about the research using this link.

3. The Devil Behind Moltbook

The research paper examines LLM-based self-evolving multi-agent AI systems and proves a key limitation in their ability to stay safely aligned while improving on their own.

The authors introduce the “Self-Evolution trilemma,” which states that an agent society cannot achieve continuous self-evolution, complete isolation, and consistent safety all at the same time.

Using an information-theoretic approach, they show that isolated self-evolution creates statistical blind spots that lead to irreversible safety lapses when compared to anthropic (human-like) value distributions.

Observations from the Moltbook agent community and other closed systems support this prediction.

The authors also claim that this loss of safety is inherent, not just a problem to fix, and stress the need for external oversight or new mechanisms to preserve safety and address these risks.

Read more about the research using this link.

4. SOTA ARC-AGI-2 Results with Agentica

The open-sourced Agentica framework by Symbolica combines code-mode agents and recursive language model (RLM) workflows using a persistent Python REPL environment.

This setup allows agents to propose, implement, test, and continuously refine programs to tackle ARC-AGI tasks more effectively than standard chain-of-thought prompting.

A significant innovation in their approach is Recursive delegation, in which agents dynamically create sub-agents to manage subtasks, allowing exploration in both depth and breadth.

Experimental results show that Agentica agents significantly improve benchmark scores for models such as Opus 4.6, GPT-5.2, and Opus 4.5, with the highest score reaching 85.28% on the public ARC-AGI-2 set.

Read more about the research using this link.

The GitHub repository associated with Agentica can be found here.

5. AttentionRetriever

Existing RAG methods struggle with long document retrieval due to poor context-awareness, causal dependence, and scope issues.

This research paper introduces AttentionRetriever, a new long-document retrieval model that shows that transformer attention layers can act as effective retrievers without any extra training.

This is achieved by using attention scores and entity-aware embeddings to capture context and scope in long texts.

Results show that AttentionRetriever can outperform previous models on long document retrieval benchmarks while being as efficient as standard dense retrievers.

Read more about the research using this link.

6. VLA-JEPA

This research paper presents VLA-JEPA, a new pretraining framework for vision-language-action (VLA) policies that improves how agents learn to act from large-scale video data.

It tackles a major issue in current latent-action objectives, which often learn irrelevant pixel changes, by using leakage-free state prediction in a latent world model where the target encoder generates latent representations from future frames, while the student pathway only sees the current observation. With such an approach, future information is used purely as a supervision target.

In contrast, by predicting in latent space rather than raw pixels, VLA-JEPA learns dynamic abstractions that are resistant to camera motion and background noise.

Experiments show that VLA-JEPA consistently improves generalization and robustness on benchmarks such as LIBERO, LIBERO-Plus, SimplerEnv, and real-world robot manipulation tasks, compared with existing methods.

Read more about the research using this link.

7. LaCy

Small language models (SLMs) have limited capacity to learn due to their small number of parameters. This often leads them to generate factually incorrect responses. To address this, SLMs are given access to external sources, such as larger models, documents, or databases, using a mechanism.

This research paper from Apple explores which tokens an SLM can and should learn during pretraining, versus which ones it should delegate via the mechanism.

Findings indicate that this choice goes beyond merely minimizing loss, and some high-loss tokens are actually reasonable alternative continuations of the training text and should be learned by the SLM instead of being delegated.

Based on this, the authors introduce LaCy, a new pretraining method that uses an augmented loss signal from the spaCy grammar parser to determine which tokens the SLM should learn to delegate to prevent factual errors and which are safe to learn and predict even under high losses.

Experiments indicate that LaCy improves factual accuracy, reflected by higher FactScores, when SLMs are combined with larger models. It also outperforms other strategies, such as Rho-trained or LLM-judge-trained SLMs, while being simpler and more cost-effective.

Read more about the research using this link.

8. Fine-Tuning GPT-5 for GPU Kernel Generation

This research paper addresses the challenge of developing efficient GPU kernels essential to high-performance AI systems, which are difficult to write because of complex hardware architectures and limited high-quality training data.

The authors show that fine-tuning GPT-5 with RL in Makora’s environment significantly improves its ability to generate correct Triton GPU kernels.

In the single-attempt setting, the fine-tuned GPT-5 improves kernel correctness from 43.7% to 77.0% and increases the fraction of problems outperforming TorchInductor from 14.8% to 21.8% compared to the baseline model, while exceeding prior state-of-the-art models on KernelBench.

When integrated into a complete coding agent, the fine-tuned model successfully solves nearly all problems in the expanded KernelBench suite and outperforms the PyTorch TorchInductor compiler on 72.9% of problems with a geometric mean speedup of 2.12X.

Read more about the research using this link.

9. GPT‑5.3‑Codex‑Spark

OpenAI has launched GPT-5.3-Codex-Spark, a research preview of an ultra-fast, real-time coding model optimized for interactive development workflows.

Spark is a smaller and highly responsive version of GPT-5.3-Codex. It is designed to deliver nearly instant code generation and edits, producing over 1,000 tokens per second with a 128k context window. This enables developers to interrupt, redirect, and improve outputs in real time.

The model runs on Cerebras’ Wafer-Scale Engine 3 hardware, making it OpenAI’s first production deployment on non-NVIDIA chips to achieve ultra-low latency.

Read more about the research using this link.

10. GLM-5

GLM-5 is Z.ai’s new flagship foundation model, built for complex systems engineering and long-horizon agentic tasks.

The model uses a Mixture-of-Experts (MoE) architecture with 744B parameters (40B active) and achieves excellent performance amongst open-source models in reasoning, coding, and agent execution.

Its capabilities are competitive with top proprietary models such as Anthropic’s Claude Opus series and OpenAI’s GPT-5.

Read more about the research using this link.

11. Gaia2

This research paper from Meta SuperIntelligence Labs introduces Gaia2, a new benchmark suite that evaluates LLM-based agents in realistic, dynamic, and asynchronous environments in which the world changes independently of the agent’s actions (unlike previous static or synchronized benchmarks).

Gaia2 tests agents on time constraints, noisy, evolving events, ambiguity resolution, and multi-agent collaboration, using a write-action verifier that supports action-level evaluation and reinforcement learning from verifiable rewards.

Initial evaluations on the benchmark show that no model performs well in all tasks. Proprietary models like GPT-5 (high) achieve the highest overall scores but struggle with tasks that are time-sensitive.

Other models vary in speed, cost, and reliability, revealing important trade-offs in reasoning, efficiency, and adaptability to real-world situations.

Gaia2 is built on a consumer environment with the open-source Agents Research Environments platform to support the development, benchmarking, and training of next-generation practical AI agent systems.

Read more about the research using this link.

This article is entirely free to read. If you loved reading it, restack and share it with others. ❤️

You can also check out my books on Gumroad and connect with me on LinkedIn to stay in touch.

Build Grouped Query Attention (GQA) From Scratch

Dr. Ashish Bamania — Fri, 13 Feb 2026 11:29:10 GMT

🎁 Here's your 25% off the annual membership. Unlock every lesson in the newsletter and accelerate your AI engineering career. Claim your discount now!

In the previous lessons on 'Into AI’, we learned to implement Self-attention and progressively updated it to reach Causal Multi-head Self-attention (MHA), which is used in LLMs like GPT-2.

MHA is a powerful architectural component that enables LLMs to process tokens in parallel, making fast autoregressive text generation possible. But it comes with high memory costs and latency.

We then learned how to build Multi-Query Attention (MQA), an efficient variant of MHA that significantly reduces memory usage and speeds up inference.

It’s now time to improve the architecture further and learn to build Grouped Query Attention (GQA).

What is Grouped Query Attention (GQA)?

In the conventional attention or Multi-Head Attention (MHA), each head has its own Q, K, and V vectors.

In Multi-Query Attention (MQA), all attention heads share the same key (K) and value (V) vectors, while each head still has its own query (Q) vector.

This reduces the KV cache size and memory requirements, leading to much faster inference.

MQA is fast and efficient, but it isn’t perfect. Because all heads share the same Keys and Values, this reduces the effectiveness of learned representations, making an LLM less expressive.

What’s the fix then?: Finding the right balance and building something that fits between MHA and MQA.

And finding that sweet spot is what led to GQA.

Grouped Query Attention (GQA) comes from a Google Research paper titled “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, published in 2023.

It has since become the de facto standard for modern LLMs and is used across all popular LLM families, including Llama, Mistral, GPT-OSS, Qwen, and more.

In GQA, Query (Q) heads are partitioned into groups, and each group shares one Key (K) and Value (V).

Grouped-Query Attention visualised (Image from author’s book titled ‘LLMs In 100 Images’)

Here is how GQA compares to MHA and MQA.

A comparison between the three:

MHA’s generation quality is the best, but its inference speed is the lowest.
MQA has the fastest inference speed and efficiency but the lowest generation quality.
GQA balances MHA’s generation quality and MQA’s inference speed/ efficiency to fit between these two.

Revisiting Multi-Head Attention (MHA)

Multi-head Attention (MHA) is implemented with causal masking in LLMs, and this architecture is called Causal Multi-head Self-attention. This prevents an LLM from peeking at future tokens while generating a token.

(Note that Causal Multi-Head Self-Attention is usually just called MHA in the popular literature.)

We have previously learned how to build Causal Multi-Head Self-Attention from scratch.

Here’s a quick recap of how it works.

The following steps are performed in this class representing MHA:

Input embeddings are projected into query (Q), key (K), and value (V) vectors
Q, K, and V are split across multiple attention heads
Scaled dot-product attention is computed, and a causal mask is applied
Softmax is applied to obtain attention weight
Attention weights are used to compute a weighted sum of values (V) for each head
The head outputs are merged and passed through a final output projection matrix

import torch
import torch.nn as nn
import math

class CausalMultiHeadSelfAttention(nn.Module):
  def __init__(self, embedding_dim, num_heads):
    super().__init__()

    # Check if embedding_dim is divisible by num_heads
    assert embedding_dim % num_heads == 0, “embedding_dim must be divisible by num_heads”

    # Embedding dimension
    self.embedding_dim = embedding_dim

    # Number of total heads
    self.num_heads = num_heads

    # Dimension of each head
    self.head_dim = embedding_dim // num_heads

    # Linear projection matrices for Q, K, V (to be split later for each head)
    self.W_q = nn.Linear(embedding_dim, embedding_dim, bias = False)
    self.W_k = nn.Linear(embedding_dim, embedding_dim, bias = False)
    self.W_v = nn.Linear(embedding_dim, embedding_dim, bias = False)

    # Linear projection matrix to produce final output
    self.W_o = nn.Linear(embedding_dim, embedding_dim, bias = False)

  def _split_heads(self, x):
    """
    Transforms input embeddings from
    [batch_size, sequence_length, embedding_dim]
    to
    [batch_size, num_heads, sequence_length, head_dim]
    """
    batch_size, sequence_length, embedding_dim = x.shape

    # Split embedding_dim into (num_heads, head_dim)
    x = x.reshape(batch_size, sequence_length, self.num_heads, self.head_dim)

    # Reorder and return the intended shape
    return x.transpose(1,2)

  def _merge_heads(self, x):
    """
    Transforms inputs from
    [batch_size, num_heads, sequence_length, head_dim]
    to
    [batch_size, sequence_length, embedding_dim]
    """
    batch_size, num_heads, sequence_length, head_dim = x.shape

    # Move sequence_length back before num_heads in the shape
    x = x.transpose(1,2)

    # Merge (num_heads, head_dim) back into embedding_dim
    embedding_dim = num_heads * head_dim
    x = x.reshape(batch_size, sequence_length, embedding_dim)

    return x

  def forward(self, x):
    batch_size, sequence_length, embedding_dim = x.shape

    # Compute Q, K, V
    Q = self.W_q(x)
    K = self.W_k(x)
    V = self.W_v(x)

    # Split them into multiple heads
    Q = self._split_heads(Q)
    K = self._split_heads(K)
    V = self._split_heads(V)

    # Calculate scaled dot-product attention
    attn_scores = Q @ K.transpose(-2, -1)

    # Scale
    attn_scores = attn_scores / math.sqrt(self.head_dim)

    # Apply causal mask (prevent attending to future positions)
    causal_mask = torch.tril(torch.ones(sequence_length, sequence_length, device=x.device)) # Create lower triangular matrix

    causal_mask = causal_mask.view(1, 1, sequence_length, sequence_length)  # Add batch and head dimensions
     
    attn_scores = attn_scores.masked_fill(causal_mask == 0, float(’-inf’)) # Mask out future positions by setting their scores to -inf

    # Apply softmax to get attention weights
    attn_weights = torch.softmax(attn_scores, dim = -1)

    # Multiply attention weights by values (V)
    weighted_values = attn_weights @ V

    # Merge head outputs
    merged_heads_output = self._merge_heads(weighted_values)

    # Final output
    output = self.W_o(merged_heads_output)

    return output

Revisiting Multi-Query Attention (MQA)

We have also learned how to implement MQA from scratch in a previous lesson.

The following steps take place in the class representing MQA:

Input embeddings are projected into query (Q), key (K), and value (V) vectors
Q is split across multiple attention heads
K and V are not split, but instead, a single shared K and V are used across all heads
Scaled dot-product attention is computed between each head’s Q and the shared K, and a causal mask is applied
Softmax is applied to obtain attention weights
Attention weights are used to compute a weighted sum of the shared values (V) for each head
The head outputs are merged and passed through a final output projection matrix

import torch
import torch.nn as nn
import math

class MultiQueryAttention(nn.Module):
  def __init__(self, embedding_dim, num_heads):
    super().__init__()

    # Check if embedding_dim is divisible by num_heads
    assert embedding_dim % num_heads == 0, "embedding_dim must be divisible by num_heads"

    # Embedding dimension
    self.embedding_dim = embedding_dim

    # Number of total heads
    self.num_heads = num_heads

    # Dimension of each head
    self.head_dim = embedding_dim // num_heads

    # Linear projection matrix for Query
    self.W_q = nn.Linear(embedding_dim, embedding_dim, bias = False)

    # Linear projection matrices for Key and Value
    self.W_k = nn.Linear(embedding_dim, self.head_dim, bias = False)
    self.W_v = nn.Linear(embedding_dim, self.head_dim, bias = False)

    # Linear projection matrix to produce final output
    self.W_o = nn.Linear(embedding_dim, embedding_dim, bias = False)

  # Splits Query into multiple heads
  def _split_heads(self, x):
    """
    Transforms input embeddings from
    [batch_size, sequence_length, embedding_dim]
    to
    [batch_size, num_heads, sequence_length, head_dim]
    """
    batch_size, sequence_length, embedding_dim = x.shape

    # Split embedding_dim into (num_heads, head_dim)
    x = x.reshape(batch_size, sequence_length, self.num_heads, self.head_dim)

    # Reorder and return the intended shape
    return x.transpose(1,2)

  # Merge heads back together
  def _merge_heads(self, x):
    """
    Transforms inputs from
    [batch_size, num_heads, sequence_length, head_dim]
    to
    [batch_size, sequence_length, embedding_dim]
    """
    batch_size, num_heads, sequence_length, head_dim = x.shape

    # Move sequence_length back before num_heads in the shape
    x = x.transpose(1,2)

    # Merge (num_heads, head_dim) back into embedding_dim
    embedding_dim = num_heads * head_dim
    x = x.reshape(batch_size, sequence_length, embedding_dim)

    return x

  # Forward pass
  def forward(self, x):
    batch_size, sequence_length, embedding_dim = x.shape

    # Compute Q, K, V
    Q = self.W_q(x) # [batch_size, sequence_length, embedding_dim]
    K = self.W_k(x) # [batch_size, sequence_length, head_dim]
    V = self.W_v(x) # [batch_size, sequence_length, head_dim]

    # Split Q into multiple heads
    Q = self._split_heads(Q) # [batch_size, num_heads, sequence_length, head_dim]

    # Add head dimension to K and V (Broadcast across all heads)
    K = K.unsqueeze(1) # [batch_size, 1, sequence_length, head_dim]
    V = V.unsqueeze(1) # [batch_size, 1, sequence_length, head_dim]

    # Calculate scaled dot-product attention
    attn_scores = Q @ K.transpose(-2, -1)
    attn_scores = attn_scores / math.sqrt(self.head_dim)

    # Create lower triangular matrix as causal masking
    causal_mask = torch.tril(torch.ones(sequence_length, sequence_length, device=x.device))

    # Add batch_size and num_heads dimensions
    causal_mask = causal_mask.view(1, 1, sequence_length, sequence_length)

    # Mask out future positions by setting their scores to -inf
    attn_scores = attn_scores.masked_fill(causal_mask == 0, float('-inf'))

    # Apply softmax to get attention weights
    attn_weights = torch.softmax(attn_scores, dim = -1)

    # Multiply attention weights by V
    weighted_values = attn_weights @ V

    # Merge head outputs
    merged_heads_output = self._merge_heads(weighted_values)

    # Obtain final output
    output = self.W_o(merged_heads_output)

    return output

Building Grouped-Query Attention (GQA)

As discussed earlier:

In GQA, Queries (Q) are partitioned into groups, and each group shares one Key (K) and Value (V).

Following this, in GQA:

We have num_heads query heads (same as MHA)
But only num_groups keys/ value heads (fewer than num_heads)
Each KV group is reused by a fixed number of query heads

Next, we will implement the GroupedQueryAttention class. This represents the causal version of GQA, which is often simply referred to as GQA in the popular literature.

We begin by setting the dimensions and creating projection matrices.

import torch 
import torch.nn as nn
import math 

class GroupedQueryAttention(nn.Module):
  def __init__(self, embedding_dim, num_heads, num_groups):
    super().__init__()

    # Check if embedding_dim is divisible by num_heads
    assert embedding_dim % num_heads == 0, "embedding_dim must be divisible by num_heads"

    # Check if num_heads is divisible by num_groups
    # (Each group must be shared by the same number of heads)
    assert num_heads % num_groups == 0, "num_heads must be divisible by num_groups"

    # Embedding dimension
    self.embedding_dim = embedding_dim

    # Number of total query heads
    self.num_heads = num_heads

    # Dimension of each head
    self.head_dim = embedding_dim // num_heads

    # Number of KV groups
    self.num_groups = num_groups

    # Number of query heads per KV group
    self.group_size = num_heads // num_groups

    # Linear projection matrix for query 
    self.W_q = nn.Linear(embedding_dim, embedding_dim, bias=False)

    # Linear projection matrices for key and value
    self.W_k = nn.Linear(embedding_dim, self.num_groups * self.head_dim, bias=False)
    self.W_v = nn.Linear(embedding_dim, self.num_groups * self.head_dim, bias=False)

    # Linear projection matrix to produce final output
    self.W_o = nn.Linear(embedding_dim, embedding_dim, bias=False)

Note how the output size for the key and value projections is num_groups * head_dim and not embedding_dim, as in the case of the query projection.

This means that we create only num_groups sets of K and V vectors per token, each shared across multiple query heads.

Next, we define a function that splits the query into multiple heads.

  # Splits Q into multiple heads
  def _split_heads(self, x):
    """
    Transforms input embeddings from
    [batch_size, sequence_length, embedding_dim]
    to
    [batch_size, num_heads, sequence_length, head_dim]
    """
    batch_size, sequence_length, embedding_dim = x.shape

    # Split embedding_dim into (num_heads, head_dim)
    x = x.reshape(batch_size, sequence_length, self.num_heads, self.head_dim)

    # Reorder and return the intended shape
    return x.transpose(1, 2)

Then, we define a function that splits K and V into num_groups heads

  # Splits K or V into num_groups heads
  def _split_groups(self, x):
    """
    Transforms K/V from
    [batch_size, sequence_length, num_groups * head_dim]
    to
    [batch_size, num_groups, sequence_length, head_dim]
    """
    batch_size, sequence_length, _ = x.shape

    x = x.reshape(batch_size, sequence_length, self.num_groups, self.head_dim)
    return x.transpose(1, 2)

Following this comes the function that merges the head outputs back.

  # Merge heads back together
  def _merge_heads(self, x):
    """
    Transforms inputs from
    [batch_size, num_heads, sequence_length, head_dim]
    to
    [batch_size, sequence_length, embedding_dim]
    """
    batch_size, num_heads, sequence_length, head_dim = x.shape

    # Move sequence_length back before num_heads in the shape
    x = x.transpose(1, 2)

    # Merge (num_heads, head_dim) back into embedding_dim
    embedding_dim = num_heads * head_dim
    x = x.reshape(batch_size, sequence_length, embedding_dim)

    return x

Finally, the forward pass is defined as follows.

  # Forward pass
  def forward(self, x):
    batch_size, sequence_length, embedding_dim = x.shape

    # Compute Q, K, V
    Q = self.W_q(x)  # [batch_size, sequence_length, embedding_dim]
    K = self.W_k(x)  # [batch_size, sequence_length, num_groups * head_dim]
    V = self.W_v(x)  # [batch_size, sequence_length, num_groups * head_dim]

    # Split Q into multiple heads
    Q = self._split_heads(Q)  # [batch_size, num_heads, sequence_length, head_dim]

    # Split K and V into num_groups heads
    K = self._split_groups(K)  # [batch_size, num_groups, sequence_length, head_dim]
    V = self._split_groups(V)  # [batch_size, num_groups, sequence_length, head_dim]

    # Expand K and V so each KV group is shared across multiple query heads
    K = K.repeat_interleave(self.group_size, dim=1)
    V = V.repeat_interleave(self.group_size, dim=1)

    # Calculate scaled dot-product attention
    attn_scores = Q @ K.transpose(-2, -1)
    attn_scores = attn_scores / math.sqrt(self.head_dim)

    # Create lower triangular matrix as causal masking
    causal_mask = torch.tril(torch.ones(sequence_length, sequence_length, device=x.device))

    # Add batch_size and num_heads dimensions
    causal_mask = causal_mask.view(1, 1, sequence_length, sequence_length)

    # Mask out future positions by setting their scores to -inf
    attn_scores = attn_scores.masked_fill(causal_mask == 0, float('-inf'))

    # Apply softmax to get attention weights
    attn_weights = torch.softmax(attn_scores, dim=-1)

    # Multiply attention weights by V
    weighted_values = attn_weights @ V  # [batch_size, num_heads, sequence_length, head_dim]

    # Merge head outputs
    merged_heads_output = self._merge_heads(weighted_values)

    # Obtain final output
    output = self.W_o(merged_heads_output)

    return output

The following operations take place during the forward pass.

Input embeddings are first projected into query (Q), key (K), and value (V) vectors
The query vector is split across multiple attention heads
The shared keys and values are split into num_groups heads
Each KV group is broadcast across multiple query heads
Each head computes scaled dot-product attention scores between Q and K
A causal mask is applied to prevent the model from attending to future tokens
Softmax converts the masked scores into attention weights
Attention weights are used to compute a weighted sum of values (V)
The outputs from all heads are merged into a single representation
A final output projection matrix produces the layer’s output

The complete code for the GroupedQueryAttention class is as follows.

import torch
import torch.nn as nn
import math

class GroupedQueryAttention(nn.Module):
  def __init__(self, embedding_dim, num_heads, num_groups):
    super().__init__()

    # Check if embedding_dim is divisible by num_heads
    assert embedding_dim % num_heads == 0, "embedding_dim must be divisible by num_heads"

    # Check if num_heads is divisible by num_groups
    # (Each group must be shared by the same number of heads)
    assert num_heads % num_groups == 0, "num_heads must be divisible by num_groups"

    # Embedding dimension
    self.embedding_dim = embedding_dim

    # Number of total query heads
    self.num_heads = num_heads

    # Dimension of each head
    self.head_dim = embedding_dim // num_heads

    # Number of KV groups
    self.num_groups = num_groups

    # Number of query heads per KV group
    self.group_size = num_heads // num_groups

    # Linear projection matrix for query
    self.W_q = nn.Linear(embedding_dim, embedding_dim, bias=False)

    # Linear projection matrices for key and value
    self.W_k = nn.Linear(embedding_dim, self.num_groups * self.head_dim, bias=False)
    self.W_v = nn.Linear(embedding_dim, self.num_groups * self.head_dim, bias=False)

    # Linear projection matrix to produce final output
    self.W_o = nn.Linear(embedding_dim, embedding_dim, bias=False)

  # Splits Q into multiple heads
  def _split_heads(self, x):
    """
    Transforms input embeddings from
    [batch_size, sequence_length, embedding_dim]
    to
    [batch_size, num_heads, sequence_length, head_dim]
    """
    batch_size, sequence_length, embedding_dim = x.shape

    # Split embedding_dim into (num_heads, head_dim)
    x = x.reshape(batch_size, sequence_length, self.num_heads, self.head_dim)

    # Reorder and return the intended shape
    return x.transpose(1, 2)

  # Splits K or V into num_groups heads
  def _split_groups(self, x):
    """
    Transforms K/V from
    [batch_size, sequence_length, num_groups * head_dim]
    to
    [batch_size, num_groups, sequence_length, head_dim]
    """
    batch_size, sequence_length, _ = x.shape

    x = x.reshape(batch_size, sequence_length, self.num_groups, self.head_dim)
    return x.transpose(1, 2)

  # Merge heads back together
  def _merge_heads(self, x):
    """
    Transforms inputs from
    [batch_size, num_heads, sequence_length, head_dim]
    to
    [batch_size, sequence_length, embedding_dim]
    """
    batch_size, num_heads, sequence_length, head_dim = x.shape

    # Move sequence_length back before num_heads in the shape
    x = x.transpose(1, 2)

    # Merge (num_heads, head_dim) back into embedding_dim
    embedding_dim = num_heads * head_dim
    x = x.reshape(batch_size, sequence_length, embedding_dim)

    return x

  # Forward pass
  def forward(self, x):
    batch_size, sequence_length, embedding_dim = x.shape

    # Compute Q, K, V
    Q = self.W_q(x)  # [batch_size, sequence_length, embedding_dim]
    K = self.W_k(x)  # [batch_size, sequence_length, num_groups * head_dim]
    V = self.W_v(x)  # [batch_size, sequence_length, num_groups * head_dim]

    # Split Q into multiple heads
    Q = self._split_heads(Q)  # [batch_size, num_heads, sequence_length, head_dim]

    # Split K and V into num_groups heads
    K = self._split_groups(K)  # [batch_size, num_groups, sequence_length, head_dim]
    V = self._split_groups(V)  # [batch_size, num_groups, sequence_length, head_dim]

    # Expand K and V so each KV group is shared across multiple query heads
    K = K.repeat_interleave(self.group_size, dim=1)
    V = V.repeat_interleave(self.group_size, dim=1)

    # Calculate scaled dot-product attention
    attn_scores = Q @ K.transpose(-2, -1)
    attn_scores = attn_scores / math.sqrt(self.head_dim)

    # Create lower triangular matrix as causal masking
    causal_mask = torch.tril(torch.ones(sequence_length, sequence_length, device=x.device))

    # Add batch_size and num_heads dimensions
    causal_mask = causal_mask.view(1, 1, sequence_length, sequence_length)

    # Mask out future positions by setting their scores to -inf
    attn_scores = attn_scores.masked_fill(causal_mask == 0, float('-inf'))

    # Apply softmax to get attention weights
    attn_weights = torch.softmax(attn_scores, dim=-1)

    # Multiply attention weights by V
    weighted_values = attn_weights @ V  # [batch_size, num_heads, sequence_length, head_dim]

    # Merge head outputs
    merged_heads_output = self._merge_heads(weighted_values)

    # Obtain final output
    output = self.W_o(merged_heads_output)

    return output

Join the Paid tier Today!

MHA, MQA & GQA Visualised

Multi-Head Attention (MHA)

MHA calculations (‘n’ in the image is the same as ‘num_heads’ in the code)

MHA weighted Values calculations (Causal mask not shown in the image)

Multi-Query Attention (MQA)

MQA operations (‘n’ in the image is the same as ‘num_heads’ in the code)

MQA weighted Values calculation (Causal mask not shown in the image)

Grouped-Query Attention (GQA)

GQA operations (‘n’ and ‘g’ in the image are the same as ‘num_heads’ and ‘num_groups’ in the code, respectively)

GQA weighted Values calculation (Causal mask not shown in the image)

Testing Out GQA

Let’s use GQA to process some randomly initialized input embeddings as follows.

# Hyperparameters
batch_size = 1
sequence_length = 4
embedding_dim = 12
num_heads = 6
num_groups = 2  # must divide num_heads

# Create input embeddings
input_embeddings = torch.rand(batch_size, sequence_length, embedding_dim)

# Initialize GQA
gqa = GroupedQueryAttention(embedding_dim, num_heads, num_groups)

# Forward pass
output = gqa(input_embeddings)

Note how GQA preserves the shape of the input ([batch_size, sequence_length, embedding_dim]) just like MHA and MQA. This allows us to stack multiple GQA layers within a Transformer block.

print("Input shape:", input_embeddings.shape)
print("Output shape:", output.shape)

"""
Output:
Input shape: torch.Size([1, 4, 12])
Output shape: torch.Size([1, 4, 12])
"""

There’s another interesting finding that you must know about:

When num_groups = 1, only 1 KV head is shared across all Q heads. This means the architecture becomes MQA.
When num_groups = num_heads, each Q head gets its own KV head. This means the architecture becomes MHA.
When 1 < num_groups < num_heads, the architecture is GQA.

If you loved reading this article and found it valuable, restack to share it with others. ❤️

If you want to get even more value from this newsletter, become a paid subscriber.

Get access to all valuable lessons, including:

Check out my books on Gumroad and connect with me on LinkedIn to stay in touch.

What is Mem0 and how does it work

Dr. Ashish Bamania — Tue, 10 Feb 2026 11:25:54 GMT

In the last part of this series, we learned about what memory is and why modern AI systems need it.

In this part, we will learn about Mem0, an open-source memory layer for LLM applications.

Become a paid subscriber to ‘Into AI’ today at a special 25% discount on the annual subscription. 🎁

What is Mem0?

🗓️ This Week In AI Research (1-7 February 26)

Dr. Ashish Bamania — Sun, 08 Feb 2026 11:38:04 GMT

Before we begin, I have an exciting offer for you that is only valid for the next 24 hours!

Grab a flat 50% discount on the annual membership of Into AI today.

👉 Upgrade here, lock in the discount, and use ‘Into AI’ as your unfair advantage:

Grab your flat 50% discount

1. GPT-5.3-Codex

OpenAI launched GPT-5.3-Codex, its most powerful agentic coding model to date.

The model combines the coding performance of GPT-5.2-Codex with the reasoning and professional knowledge of GPT-5.2 in one system. It runs about 25% faster than the previous version and excels in industry benchmarks, such as SWE-Bench Pro and Terminal-Bench.

It can handle long, multi-step tasks that involve research, tool use, debugging, deployment, and complex execution.

And unlike the earlier versions, GPT-5.3-Codex not only writes and reviews code but also acts as an interactive collaborator throughout the software lifecycle and other knowledge work.

One of the most interesting aspects is that the Codex team used early versions of GPT-5.3-Codex to debug its own training, manage its own deployment, and analyze test results and evaluations, making it the first OpenAI model to help create itself.

Read more about this research using this link.

2. Claude Opus 4.6

Anthropic has released Claude Opus 4.6, its most capable model to date.

Opus 4.6 performs exceptionally well in coding, planning, and long-term tasks, and is more reliable with large codebases and complex workflows, with a stronger ability to find and fix its own mistakes.

Additionally, it introduces a 1-million-token context window (in beta), which improves its understanding of longer contexts.

On benchmarks, Opus 4.6:

Scores 2nd on Terminal-Bench 2.0 for agentic coding, only next to GPT-5.3-Codex

Tops Humanity’s Last Exam for complex reasoning
Performs best on GDPval-AA for economically valuable knowledge work
Ranks highest on BrowseComp for finding hard-to-locate information online

Opus 4.6 can assist with research and financial analysis, create documents, work with spreadsheets and presentations, and multitask autonomously.

Despite these advancements, it maintains strong safety and alignment standards that meet the highest industry standards.

Read more about this research using this link.

3. Qwen3-Coder-Next

Alibaba’s Qwen team introduced its new open-weight AI coding model called Qwen3-Coder-Next, built specifically for agentic coding tasks and local development workflows.

This model uses an 80B-parameter Mixture-of-Experts (MoE) architecture and activates only about 3B parameters during inference, making it efficient and cost-effective.

Despite using so few parameters, it delivers high performance on real-world coding challenges, including long-horizon reasoning, complex tool use, and recovery from execution failures.

Read more about this research using this link.

4. OpenAI Frontier

OpenAI has launched OpenAI Frontier, a new platform for businesses to build, deploy, manage, and govern AI agents.

This platform connects agents with company data and tools, such as CRM systems, data warehouses, and internal apps, and provides shared context, onboarding, feedback loops, permissions, and governance.

It also allows agents to automate complex tasks, including data analysis, software engineering, and operations.

Frontier supports both agents built on OpenAI technology and third-party agents, aims to reduce fragmentation (Agent sprawl), and helps companies use AI effectively within their existing systems.

5. Voxtral-transcribe-2

Mistral AI released Voxtral Transcribe 2 family of speech-to-text models that includes the following:

Voxtral Mini Transcribe V2: For batch audio transcription
Voxtral Realtime: For live, ultra-low latency streaming (that transcribes at the speed of sound)

Voxtral Mini Transcribe V2 is a batch speech-to-text model designed to transcribe pre-recorded audio, with a focus on accuracy and a detailed transcript structure.

It includes speaker diarization (who spoke when), word-level timestamps, strong noise robustness, and context biasing (better handling of names and domain jargon).

Voxtral Mini Transcribe V2 achieves the lowest word error rate at the lowest price point in the industry ($0.003/min).

Voxtral Realtime is a live, streaming speech-to-text model, designed for ultra-low-latency transcription (sub-200 ms).

This makes it suitable for real-time voice agents and for supporting key features such as timestamps, speaker separation, and stability in noisy real-world audio environments.

Voxtral Realtime is available as open source under the Apache 2.0 license and is efficient enough to be deployable on edge devices for privacy-first applications.

6. World Action Models are Zero-shot Policies

This research by NVIDIA presents DreamZero, a 14B-parameter World Action Model (WAM) built upon a pretrained video diffusion backbone and designed for robot control.

DreamZero learns to predict future world states and corresponding actions as part of the same model. This allows it to generalize zero-shot to new tasks and environments better than existing VLA models, which struggle with unseen physical motions (over 2× improvement in task progress in unseen scenarios).

It works in real time and can quickly adapt to a new robot using about 30 minutes of play data. This demonstrates its strong cross-embodiment transfer and robustness without needing repetitive demonstrations.

Read more about this research using this link.

7. Kona 1.0

Logical Intelligence presents an Energy-Based Reasoning Model (EBRM) called Kona 1.0 that can solve hard Sudoku puzzles, showing that EBMs can handle constraint-satisfaction problems much better than standard LLMs.

Kona does not generate solutions one token at a time. Instead, it evaluates entire candidate solutions against all constraints at once, assigning a low energy score to valid grids and using gradient-based adjustments to fix any violations.

In their benchmark, Kona solved 96.2% of puzzles in about 313 milliseconds, while leading LLMs like GPT-5.2, Claude Opus 4.5, and Gemini 3 Pro only solved around 2% before timing out or failing.

This shows that LLMs’ autoregressive design makes them poorly suited to reasoning on spatial tasks, whereas energy-based models can reason holistically and efficiently revise partial solutions to these problems.

Read more about this research using this link.

Try out the demo yourself using this link.

Grab your flat 50% discount

8. Mining Generalizable Activation Function

This research from Google DeepMind presents four novel activation functions discovered using an evolutionary search powered by AlphaEvolve.

These activation functions are:

GELU-Sinc-Perturbation
GELUSine
Gaussian-Modulated Tangent Unit (GMTU)
Turbulent Activation Function

These activations consistently outperform ReLU and often GELU, especially on out-of-distribution (OOD) and reasoning-style benchmarks.

Read more about this research using this link.

9. Learning to Reason in 13 Parameters

This research from Meta presents TinyLoRA, an extremely low-parameter fine-tuning method for LLMs that enables reasoning with very few trained parameters.

The authors show that applying TinyLoRA to a base model (Qwen-2.5-7B-Instruct) can significantly improve performance on reasoning benchmarks such as GSM8K, achieving about 91% accuracy while training with only 13 parameters.

This trend continues with tougher learning-to-reason benchmarks, including AIME, AMC, and MATH500, recovering around 90% of the full performance gains while using 1000X fewer trainable parameters.

Note that such strong performance was achieved only with RL training. Models trained using SFT required 100-1000x more trainable parameters to achieve the same performance.

Read more about this research using this link.

10. EchoJEPA: A Latent Predictive Foundation Model for Echocardiography

This research paper introduces EchoJEPA, a new latent predictive foundation model trained on 18 million echocardiograms from about 300,000 patients, the largest dataset for this type of imaging to date.

The model uses a latent-prediction objective to learn strong anatomical representations that ignore speckle noise and artifacts inherent to ultrasound, thereby improving clinical metrics.

Compared with leading baselines, EchoJEPA achieves about 20% better performance in estimating left ventricular ejection fraction (LVEF) and 17% better performance in estimating right ventricular systolic pressure (RVSP).

Furthermore, it generalizes significantly better under acoustic disturbances and performs well on pediatric data without any fine-tuning.

The research suggests that latent predictive modeling is a promising approach for developing reliable and generalizable medical imaging AI.

Read more about this research using this link.

This article is entirely free to read. If you loved reading it, restack and share it with others. ❤️

Don’t forget to grab your flat 50% discount on the annual membership of Into AI today.

👉 Upgrade here, lock in the discount, and use ‘Into AI’ as your unfair advantage:

You can also check out my books on Gumroad and connect with me on LinkedIn to stay in touch.

How RNNs Work (And Why Everyone Stopped Using Them)

Dr. Ashish Bamania — Sat, 07 Feb 2026 11:00:51 GMT

Reshape (Affiliate Partner)

Tired of trying to reach your weight loss goals and failing each time? We have all been there. But not anymore.

It’s time to start seeing real results with Reshape!

Just snap a photo or say what you ate, and instantly get precise calorie, macro, and nutrient tracking with Reshape’s AI, skipping all that tedious logging work.
Get workout plans tailored for your equipment, schedule, and fitness level that change as you progress. Whether you’re building muscle, losing fat, or getting stronger, these plans adapt to your needs.
Your 24/7 AI coach, Fio, is a personal trainer in your pocket that provides instant feedback on your form, meal ideas, and motivation whenever you need it.
Track everything that matters in one powerful dashboard: meals, workouts, sleep, steps, and body composition, with easy Apple Health integration.
Track your weekly insights to learn what’s working for you so that you keep winning each and every day.

Use the code “INTOAI” for a free Reshape Pro trial and join 45K+ people who have transformed their bodies with Reshape. Your smarter fitness journey starts now!

Try Reshape for FREE!

This week’s newsletter is written by . He is a senior Data Science manager at Skyscanner.

He regularly shares insights on building effective teams, developing leadership skills, and advancing careers in Data Science and Machine Learning through his newsletter, Senior Data Science Lead.

You can also find him and stay up to date with his content on LinkedIn.

This week, I went all the way back to 2017. That’s when the now-legendary ‘Attention is All You Need’ paper came out — the one that introduced the world to Transformers, and set the foundation for everything from ChatGPT to image generation to code-writing copilots.

And sure, I could jump straight into explaining how Transformers work. But given the impact these models have had — and the fact that you probably see the word “attention” 30 times a week now — I thought it would be worth taking a step back (actually 2 steps back).

Before we can truly understand Transformers, we need to understand where they came from. And that means revisiting the architectures that paved the way: Recurrent Neural Networks (RNNs) and Long Short-Term Memory (LSTM) networks.

In this blog post, I am diving into RNNs.

We will walk through what they are, how they work, and most importantly, why they struggle. By the end, RNNs will feel like a clever little for-loop with memory instead of scary maths magic.

Ready? Let’s jump in!

What we will cover in this article

How is an RNN different from a classical Deep Neural Network? (or other classical sequence models)
Introducing a made-up use case with 3 data points for Stock price prediction
The scary official diagram of an RNN cell. (Don’t worry, we will break it down super easily.)
A walkthrough of RNN calculations (And you will see how the maths is not that scary after all.)
Three problems associated with RNNs

How is an RNN different from a classical Deep Neural Network?

I assume in this post that you have worked with (or are familiar with) the basics of classical Deep Neural Networks (from here on, DNNs).

If my assumption is correct, then the diagram below should feel really familiar. It is a diagram of a DNN with:

X: A set of input nodes, representing the variables you want to use for prediction.
H1, H2: 2 hidden layers with 4 nodes. This is where the parameters that learn how to dial up or down specific signals from X. Basically, the knobs the model adjusts to learn.
Y: The prediction node. In this case, it’s only 1 for simplicity.

Diagram of a classical Deep neural network (DNN)

Now, there are 2 main things to highlight in this diagram:

The DNN processes all the data at once.
You see the X input? From a DNN’s perspective, it’s a torrent of data, all pushed and processed at once. There is nowhere in this diagram that the network can say, “Hey, can you just send me X1 and X2 first, and then I can process X3 and X4?”.
The DNN is feed-forward (or sequential).
In other words, the data flows from left to right (from X → Y). There is nowhere in this diagram that a node can stop the data flow and ask: “Hey, what data did I have in the previous step?” It is oblivious to that.

DNNs are really powerful, but they are also “memory-less”

DNNs are “memory-less” because of the two points mentioned above. And being “memory-less” means that DNNs really struggle to predict when sequence or order matters.

Take a simple stock price prediction scenario. The only way that DNNs can consider what happened yesterday or the day before is if you tell them what matters. This is usually done by manually creating features like:

Yesterday’s price
A moving average over the past 7 days
A sine-transformed day-of-week feature to capture seasonality

But wasn’t Deep learning supposed to eliminate Feature engineering?

Yes (and no).

Neural networks do learn internal representations (i.e., “features”) from raw data. But when it comes to sequences, classical deep neural nets still need all the help they can get (so, kind of back again to square one, where we have to feature engineer stuff…)

This is where Recurrent Neural Networks (RNNs) come in.

Recurrent neural networks (RNNs) are a class of artificial neural networks designed for processing sequential data, such as text, speech, and time series, where the order of elements is important.

Their two main characteristics are:

RNNs process one input at a time.
Instead of taking all of the input X in one big gulp, RNNs look at one data point at a time — like reading a sentence word by word, or stepping through a stock price day by day. This allows them to focus on how each input evolves over time.
RNNs are recursive.
Yes, the information flows left to right, but at each step, it can also look at what happened before (kind of a right-to-left motion). It’s like a left-to-right with memory.

Don’t worry if this feels dense right now, we will break it down step by step.

By the end of this post, you will not only understand what “recurrent connections” mean, but you will also see why RNNs became a foundational architecture for handling sequences.

Introducing a made-up use case for Stock price prediction

Ok, before we get into the RNN section, let me introduce you to the simplest stock price prediction exercise in the course of human history.

The diagram below shows a toy time series where:

For simplicity, all stock prices are set to 0.
We have 3 data points.
The idea is to use yesterday’s and today’s stock prices to predict tomorrow’s stock price.

I want to introduce this in such a simple way because:

I want to actually show you the maths of the RNN using super simple numbers.
Labelling each step as yesterday, today, and tomorrow will help anchor the RNN diagrams that follow.

Let’s keep this mental model in our back pocket — it’s going to make the scary-looking RNN diagrams feel a lot less scary.

Ok, now we are ready to get scared by a diagram of an RNN cell!

The scary official diagram of an RNN cell

So far, we have discussed how RNNs differ from classical DNNs because they remember their past. But what does that actually look like inside the model?

Well... time to face the infamous RNN diagram.

It might look like a tangle of wires and equations at first, but don’t worry, we will walk through it slowly, tie it back to the stock price example, and by the end, you will see it’s just a simple process of multiplication, addition, and a squiggly activation function or two.

The diagram below is a vanilla RNN cell.

I can sense you sweating… a flow diagram? With parameters? With maths operations? Ok, let’s break it up so that you don’t have to process it all at once.

Here are the elements to focus on:

Note the X, h and Y elements (similar to the classical DNN diagram).
X(t) represents today's data point, h(t-1) is what came from yesterday, h(t) is what is passed to predict tomorrow, and Y(t+1) is the predicted data point.
Y(t+1) represents the prediction that we want.
tanh and softmax are just activation functions. The same kind you have seen in regular neural networks, so nothing special for RNNs. They take raw values and squash them into a friendlier range, like between -1 and 1 (for tanh) or 0 and 1 (for softmax).
Finally, there are a couple of blobs with mult and sum. These are just visual aids so that you can see the operations in action when we pull the numbers in.
There are maths operations outside of the cell. These are basically there just so that we can transform a squashed value coming out of tanh, into a real value that makes sense. For example, transform 0.9 coming out of the tanh function to maybe $5.

✍️ Quick note

Technically, the RNN makes a prediction y(t) after seeing input x(t) and memory h(t-1).

But since in our toy example we are trying to predict the next value, it’s tempting to call it y(t+1) — just know that it’s really y(t) in the math, but the target we are aiming for is the value at t+1.
I purposely named it y(t+1) for pedagogical reasons for this post.

Let’s map this diagram with the theoretical RNN math function

In the image below, I have added the math functions that lead to the two outputs from the RNN cell: y(t+1) and h(t).

So, if we wanted to either predict an output (Y) or carry to the next stage (h), then the neural network should learn:

Wx: This is the weight applied to today’s data, X(t). It controls how much the model should care about today’s stock price. Extreme case, if Wx = 0, then this means we don’t care about X(t) because Wx * X(t) would yield 0.
Wh: This is the weight applied to the previous hidden state, h(t–1). It tells the model how much to rely on memory. If Wh = 0, the past is forgotten.
⚠️ Don’t misinterpret h(t-1) by thinking it is yesterday’s stock price. It is what comes out of the cell (with its multiplications, sums, and activation functions) applied to yesterday’s stock price.
b: This is the bias term. Think of it as a small correction applied regardless of the input. It’s important in training, but not very interesting for understanding how RNNs work conceptually (as it affects DNNs the same way). If you want to deep dive, check this link.

I believe that only when we plug in numbers to these diagrams will we start really understanding what is happening inside the RNN. Let’s do this next.

A Walkthrough With Real Numbers

Alright, time to take what we have learnt and run it step by step. Instead of just showing the internals of a single RNN cell, we’ll now “unroll” it (you will see what that is in a second).

This will finally answer the big question: how does an RNN actually use the past to predict the future?

How do we represent tomorrow’s prediction diagrammatically?

Pretty simple. We just copy and paste the same RNN cell forward in time, one per data point in our sequence.

In our toy stock price example, we only have two data points (t–1 and t), so we unroll the cell twice in order to make a prediction at t+1.

And this is why they are called recurrent, because the same logic is applied over and over, like a for loop.

Cool, now that you are comfortable with what happens inside a single RNN cell, let’s walk through this unrolled diagram in detail to ensure we are all on the same page.

We begin by plugging in yesterday’s stock price. Because there is no data prior to yesterday, we can ignore the previous hidden state input h(t-2).
Using both math functions, we then calculate y(t) and h(t-1). From these two, only h(t-1) is useful for us. This is the value that describes the memory, and that will be passed to the next cell. y(t) is irrelevant, so we ignore it.
Finally, we plug in today’s data and run through the relevant maths operations to calculate y(t+1). You can see from the diagram that the RNN is using today’s data X(t) and yesterday’s data h(t-1) from memory to calculate what could happen tomorrow.

💡 An important highlight: Shared weights and biases

“Wait a sec... you are using the same Wx, Wh, and b in both cells. Shouldn’t they be different?”

Great question! This is the part that makes RNNs elegant, but also tricky (you will see at the end how these shared weights break an RNN’s learning process).

Unlike feedforward layers that might learn new weights for every input, an RNN cell reuses the same weights at every timestep. That’s the “recurrent” part. An RNN not only repeats the cell structure, but it also repeats the exact same function with the same learned parameters.

So yes — Wx, Wh, and b are constant across time. What changes is the input x(t) and the memory from the previous step h(t–1), which is how the model updates its thinking as it moves forward.

Plugging numbers into the diagram

Before doing some basic maths, let’s talk about the numbers being used:

Note how I substituted Wh, Wx, b, W_output and b_output with numbers. I made these numbers up, but they are the ones the neural network would tweak during its learning process.
The input data points are the ones we know from the time series. X(t-1) and X(t) are both 0.

Now we are ready to take pen and paper and perform all the calculations in this diagram.

Plugging in all the numbers, y(t+1) comes out to be 0. Nice, this is what we expected from our mock stock price time series!

What would the diagram look like if we had 50 data points?

Well... you would copy the RNN cell forward 49 times, just like a for loop with 50 iterations (luckily for me, I am not drawing that diagram...)

But structurally, nothing changes. You still:

Reuse the same weights
Pass memory from each time step to the next
And only apply the final output prediction where it matters, which is usually the last cell in the sequence

What we learn from the above examples is that:

RNNs reuse the same weights and biases at each step.
Hidden states h(t) are like memory, passed forward through time.
Outputs y(t) are generated by applying a linear transformation and softmax on the hidden state.
Even a dumb toy dataset of all zeros reveals the internal mechanics beautifully.

That wasn’t that hard, right? Well, I have bad news for you…

Three Major Problems With RNNs

These vanilla RNNs are never used for real-world use cases because they come with three big problems:

They are slow to train.
They suffer from the problem of vanishing gradients
They suffer from the problem of the exploding gradients

Let’s cover these in detail in their own sections.

Problem 1: Training a vanilla RNN is very slow

Deep Neural Networks (DNNs) process all their inputs at once in a single forward pass. Everything flows from left to right, layer by layer. That means training can occur in parallel across many data points and GPU cores.

This is not the case with RNNs.

Because RNNs depend on previous hidden states, they are inherently sequential. You can’t calculate h(t) until you have calculated h(t–1).

It’s like reading a book: you can’t understand chapter 5 until you have read chapter 4.

RNNs make you walk through time, one step at a time, and this sequential dependency kills parallelism.

That is the first problem: RNNs are powerful, but they pay for it in training speed.

Problem 2: The problem of the Vanishing gradients

To explain this problem, I would have to take some mathematical shortcuts. To really know what is happening under the hood, you need to be familiar with the chain rule used in backpropagation. But showing the full impact of vanishing gradients using the backpropagation formula would be overkill.

Let’s zoom in on just one parameter in the RNN, i.e. Wh, the weight that is multiplied by the previous hidden state.

In our earlier 3-step example, Wh was used once (just one multiplication). But, if you are training an RNN for 50 timesteps (say, 50 days of stock data), that means Wh shows up 49 times in the full chain of calculations.

When the model tries to update Wh via backpropagation, mathematically, the gradient is multiplying Wh over and over again, kind of like:

Uh-oh. What happens when Wh < 1?

Try plugging in Wh = 0.5:

That’s basically zero. This means that the update to Wh during training is so tiny, it’s like the model is frozen. It can’t escape its starting point. It just sits there, unable to learn anything useful.

Problem 3: The problem of the Exploding gradients

This is the opposite problem. If vanishing gradients are the slow death of learning, exploding gradients are the chaotic opposite.

What happens when Wh > 1? Say Wh = 1.5.

That’s nearly a billion. This means that during backpropagation, the gradient becomes massive. And with a gradient that large, your weight update becomes a wild jump.

The result is that your model overshoots the loss minimum, bounces around the optimisation landscape like a drunk pinball (love that game), and probably never converges.

This is the exploding gradients problem. Same root cause as vanishing gradients — compounding multiplications through time — but now the problem is too much signal, instead of too little.

What’s next: How LSTMs fixed all of this (mostly)

With everything we have covered, you might think, “Hey, that vanilla RNN was easy enough to understand, but with the stated problems, it also looks pretty useless, right?”

I would mostly agree. Training a vanilla RNN is not impossible, but it requires skill, feature engineering, fine-tuning, and time. The vanishing and exploding gradient problem is the one that mostly holds back an RNN.

This is why LSTMs were introduced.

LSTMs are improved versions of RNNs. With built-in mechanisms (called Gates) to decide what to keep, what to forget, and what to pass forward, they were designed specifically to beat vanishing and exploding gradients at their own game.

LSTMs are for the next post, where we will explore how they work and why they became the go-to tool for sequence modelling… at least until Transformers came along.

Now, I want to hear from you!

In this post, we broke down how RNNs work, from the vanilla cell structure to why they struggle to train in the real world.

We kept it simple, even used an all-zero dataset, and uncovered the quiet math that makes RNNs nearly impossible to scale.

But now I’m curious about your experience.

Have you ever built or trained an RNN model?
Did you run into vanishing gradients or the joys of exploding updates?
Maybe you were introduced to LSTMs (or GRUs) straight away and skipped vanilla RNNs altogether?
Or maybe you are just now connecting the dots between hidden states, time steps, and why Transformers were such a leap.

Drop your thoughts, experiences, or lingering questions in the comments. I would love to hear how you’ve approached sequence modelling in your ML journey.

Thanks again to for writing this week’s newsletter.

Don’t forget to subscribe to his newsletter and connect with him on LinkedIn.

This article is free to read, so share it with others. ❤️

If you want to get even more value from this publication and support me in creating these in-depth tutorials, consider becoming a paid subscriber.

You can also check out my books on Gumroad and connect with me on LinkedIn to stay in touch.

Memory For AI Agents: Everything That You Need To Know (Part-1)

Dr. Ashish Bamania — Wed, 04 Feb 2026 10:30:36 GMT

🎁 Become a paid subscriber to ‘Into AI’ today at a special 25% discount on the annual subscription.

LLMs are stateless. This means that their state is not permanent and refreshes with each session. This is why LLMs do not remember past interactions between sessions.

AI Memory systems aim to fix this problem. They help LLMs retain and use past information. This helps them stay relevant, learn over time, and respond more consistently in user interactions, grounded in truth.

What are the different types of Memory?

Modern-day AI memory systems are based on how humans store and retrieve knowledge.

The Soar cognitive architecture is one of the earliest and most comprehensive attempts to model human-like intelligence in agents (in our case, an LLM).

According to it, there are two types of memory:

Working memory, or Short-term memory, which holds information about what the agent is thinking about during ongoing problem-solving.
Long-term memory, which is further divided into three:
- Procedural memory: This stores rules or skills that can be applied to working memory to determine the agent’s behavior.
- Semantic memory: This stores general facts and knowledge about the world.
- Episodic memory: This stores sequences of the agent’s past behaviors and interactions.

Learn about this architecture makes discussing memory in LLMs easier.

Are LLMs completely memoryless?

It’s not that LLMs lack all sorts of memory by default.

LLMs have a context window that stores information about an ongoing conversation with a user/ another agent. The context window is the maximum length of input an LLM can consider at once.

This acts as an LLM’s Short-term or Working memory.

The context window is a limited space for storing short-term information, but recent LLMs have really pushed its boundaries.

LLMs like Google’s Gemini 2.5 and xAI’s Grok 3 have a context window of 1 million tokens, which means that they can process up to 1,500 pages of text or 30K lines of code simultaneously!

Context Window (Image from LLMs In 100 Images)

Alongside this, LLMs have something called Parametric memory, which is the knowledge that is stored directly in their parameters / weights as a result of training.

For example, if you ask an LLM about the capital of the UK, it will respond with London.

This is an LLM’s Long-term memory (specifically Procedural and Semantic memory).

Note that Procedural memory also includes an LLM’s system prompt that persists across sessions.

If LLMs already have memory, why do we need more?

Sure, LLMs can use their context window as their working/ short-term memory, but this is limited.

Meaningful human interactions and relationship-building take place over a long period (weeks, months, or years). As conversations accumulate and get longer, the context window fills up pretty quickly.

There is also no thematic continuity in conversations, which means a person could talk about writing code one day, then describe how sad and lonely they feel about a past relationship, and then return to writing code another day.

Even if we capture this interaction information in the context window, store it in a persistent database (creating Long-term memory), and reload it each time, such a discontinuity in the conversational theme in a huge conversation trail creates a messy context (a phenomenon called Context Rot).

Even with a million-token context window, LLMs cannot always accurately extract the most meaningful information from such a context in their context window.

Check out the Needle-in-a-Haystack (NIAH) analysis of GPT-4, which checks whether the model can find a small, specific fact (‘Needle’) hidden somewhere in a very long context (‘Haystack’).

GPT-4 performs very poorly on retrieval accuracy as context length increases.

Many modern-day LLMs have significantly improved on the NIAH test and show nearly perfect accuracy, but they still perform poorly on more comprehensive benchmarks like RULER.

RULER checks how well an LLM can really use a long text context (information tracking, aggregating, and reasoning across long contexts), and not just find bits of information in it.

Claimed context length and effective context length of different LLMs based on RULER (Source)

It seems that pushing the limits of context length only delays, rather than resolves, the fundamental limitation.

That’s about the problems with LLM’s current short-term memory and using it towards building long-term memory.

There’s another way where we can directly update an LLM's long-term memory. This is with mid-training and post-training methods, but this is an expensive process (in terms of time, money, and compute).

And even if we try to do this, there’s the curse of Catastrophic forgetting, where LLMs can abruptly and drastically forget previously learned information upon being trained with new information.

This is why AI needs memory that goes beyond just longer context windows.

A robust AI memory should:

Selectively store important information
Connect related information
Provide the right information when needed
Forget information that is not frequently needed over a long period

much like how humans remember things.

With this type of memory, an LLM-based AI system can stay consistent, remember user preferences, build on past conversations, and become a reliable long-term partner that grows with the user.

That’s everything for the first part of this series on AI memory.

In the next one, we will discuss a memory system called Mem0 in depth, which acts as a universal, self‑improving memory layer for LLM applications.

Upgrade to a paid subscription today and use ‘Into AI’ as your unfair advantage: