This article is also available as a podcast! If you’re on the go or just want to absorb the content in audio format, you can listen to the full episode below 👇 The podcast is also available on Spotify and Apple Podcasts.

As LLM workflows move from simple chat experiments to production-grade automation, developers are starting to hit a “prompt ceiling.” Traditional prompt engineering relies on large system prompts that are difficult to version, hard to maintain across teams, and increasingly fragile as context windows grow. Instructions must often be repeated, and small wording changes can significantly alter model behavior.

Until recently, the main alternatives were either platform-locked assistants inside apps like ChatGPT or the Gemini app (e.g., custom GPTs or Gems) or more complex infrastructures built around function calling and custom orchestration. Neither approach provides a simple way to package and reuse procedural knowledge across projects.

A more structured alternative has recently emerged with Claude Skills, introduced by Anthropic. Instead of embedding workflows inside prompts, developers can package expertise into skills: structured folders containing instructions, resources, and scripts that Claude can discover and load dynamically when needed. This approach turns AI workflows into modular, versionable components that can live in a filesystem or repository alongside the rest of a project.

Objective

This article focuses on Agent Skills for Claude: what they are, how they differ from existing approaches, and how to build one from scratch.

After reading this article, you will understand:

1. Why prompt engineering does not scale, and the limitations shared by existing alternatives

2. How Agent Skills work, from their internal structure to their discovery and loading model

3. How to design, iterate, and deploy a Claude Skill in practice across Claude’s different surfaces

Prerequisites — None. This article is self-contained and assumes no prior knowledge of Agent Skills or Claude Code.

Subscribe now

1. The Limits of Current LLM Workflow Approaches

1.1 The Limits of Prompt-Based Workflows

Why Traditional Prompting Does Not Scale

Prompt engineering, which consists of crafting inputs to guide model outputs, became the dominant way to control LLMs without modifying their weights. It is attractive because it requires no infrastructure or training pipeline: just text.

However, control expressed through text is fragile. Instructions must be repeated in every session, degrade as context grows, and are difficult to version or reuse. Prompts are also model-sensitive: one tuned for GPT-4o may behave differently on Claude 3.5 Sonnet or even on a newer model version.

As a result, over 65% of LLM developers report that prompt versioning and observability are among the main challenges when moving from prototypes to production [1].

Repeated Instructions and Fragile Prompt Logic

Unlike traditional software, which runs deterministic code, prompt engineering operates in a probabilistic system. Small wording changes can produce very different outputs, and identical prompts may yield different results across sessions, making reproducibility difficult.

Teams often end up copying large system prompts across projects, creating diverging versions that are hard to maintain. Over time, prompts grow into long instruction blocks where even small edits risk breaking hidden logic [1].

The Challenge of Reusable AI Workflows

The deeper issue is knowledge portability. Expertise often lives in documents or personal notes rather than in a structured form that models can reliably access.

Prompt engineering works well for experimentation, but it is a weak foundation for building reusable and maintainable AI workflows.

1.2 Existing Approaches to Extending LLMs

Before Skills, several strategies emerged to make LLMs more capable or specialized.

Custom GPTs, Gemini Gems, and ChatGPT / Claude Projects

These are the closest alternatives. Each lets you create a persistent assistant with instructions and optional knowledge files, without additional infrastructure. They work well for interactive chat use, but have clear limits for developer workflows.

1. Platform lock-in: these configurations live inside their respective interfaces and cannot be version-controlled, packaged, or used programmatically.

2. Monolithic by design: each is a single, fixed configuration. If you need slide generation and SQL generation and report formatting, you maintain separate assistants with no way to combine them dynamically in a single session.

3. Static knowledge, no execution: uploaded files are read-only context, not code LLM can run. Anything that cannot be expressed in prose is simply out of scope.

Function Calling

Function calling gives the model the ability to act, calling APIs, writing files, or querying databases. But it does not define behavioral patterns or workflows across sessions.

2. From Prompts to Skills: A New Interaction Model for LLMs

2.1 Skills: Packaging Expertise Instead of Writing Prompts

Introduced by Anthropic in October 2025, Agent Skills represent a new way to structure the way language models perform tasks. Instead of repeatedly crafting prompts, developers can package procedural knowledge into reusable capabilities called skills.

A skill is essentially a structured folder containing instructions, scripts, and supporting resources that an agent can discover and load dynamically. Once created, it can be reused across conversations and projects without rewriting prompts.

In other words, building a skill is like writing an onboarding guide for a new employee. Procedural expertise is documented once in a structured format, and the model applies it whenever relevant.

In December 2025, Anthropic released Agent Skills as an open standard, signaling a broader ambition: making skills portable across tools and platforms rather than tied to a single product ecosystem.

For teams, the implications are significant. Skills can be stored in a filesystem or Git repository, versioned alongside code, and shared across projects. This makes LLM workflows into maintainable infrastructure rather than prompt engineering experiments.

2.2 Anatomy of a Skill

At the most basic level, a skill requires only one file: SKILL.md. Everything else in the directory is optional and loaded only when needed.

To illustrate the structure, consider a skill designed to generate branded client presentations from predefined templates and upload them to Google Drive:

branded-slides/
├── SKILL.md
├── assets/
│   ├── infographics-library.pptx
│   └── slide-example.pptx
├── references/
│   ├── brand-guidelines.md
│   └── google-slides-setup.md
└── scripts/
    ├── credentials.json
    └── upload_to_google_slides.py

The SKILL.md file begins with YAML frontmatter, which tells the model when the skill should be activated.

---
name: branded-slides
description: Generate branded client slide decks as .pptx files and upload
  to Google Drive. Use when asked to create a presentation, deck, or slides
  for a client or course.
---

Only two fields are required:

• name : unique identifier for the skill

• description : the trigger description the model uses to decide whether to activate the skill

Because the model relies entirely on the description to detect when to invoke the skill, specific wording matters: a vague description reduces the likelihood that the skill will be triggered automatically.

Skills can optionally declare the tools they are allowed to use via the allowed-tools field:

---
name: branded-slides
description: Generate and upload branded slide decks to Google Drive.
allowed-tools: Bash, Read, Write
---

In practice, the tool scope should match the risk level of the workflow. For example:

Example 1 — Read-only skill (safe analysis)

A skill designed only to inspect files should not have permission to modify them.

# Read-only skill — analysis without any risk of modification
allowed-tools: Read, Grep, Glob

This configuration allows the model to search and analyze a codebase or document collection without any possibility of altering files.

Example 2 — Full file workflow

More complex automation tasks require the ability to read files, generate content, and write results.

# Full file workflow — read, generate, write
allowed-tools: Bash, Read, Write

This setup is typical for automation workflows, such as generating reports, building artifacts, or running scripts that produce output files.

Below the frontmatter, the rest of SKILL.md contains the procedural workflow the model should follow.

# Branded Slides

## Instructions
1. Read references/brand-guidelines.md
2. Inspect assets/infographics-library.pptx
3. Run scripts/upload_to_google_slides.py

## Constraints
- Start with a title slide and executive summary
- Use amber #FBB831 for headers
- Maximum 5 bullet points per slide

Supporting files play distinct roles:

• brand-guidelines.md defines visual and branding rules

• upload_to_google_slides.py performs deterministic operations like generating slides and uploading them

• the assets/ folder provides visual templates and references

2.3 Progressive Loading & Skill Discovery

A key architectural feature of skills is progressive disclosure. Instead of loading all information at once, the model retrieves instructions only when needed.

Information is loaded in three stages:

Skills are invoked by the model, not the user. The agent decides when to activate a skill based on the request and the skill’s description. Initially, the model only knows the skill exists. When a relevant request appears, it loads the instructions and fetches additional resources only if required.

Unlike system prompts, which are always present and consume tokens, skills stay dormant until needed and load only the necessary information.

This modular approach lets teams treat LLM capabilities as reusable building blocks. For example, a team might create separate skills for generating slide decks, analyzing SQL queries, or producing reports, and the model activates each one only when the task requires it.

3. Building a Claude Skill in Practice

3.1 Start With the Workflow & Use Claude as the Skill Writer

The most effective way to build a skill is not to start with a text editor. Instead, begin in Claude and work through the task in a normal conversation.

For example, when building a skill to generate branded client presentations, the first step is simply asking Claude to produce a presentation and guiding the process: providing brand rules, clarifying slide structure, correcting formatting, and adding any steps needed to produce the final output.

This approach works because the conversation naturally contains everything a skill requires: brand constraints, file paths, sequencing rules, and edge cases. The workflow itself becomes the raw material.

Once the task is complete, a single prompt is enough:

“Create a SKILL.md from the workflow we just used.”

Claude understands the Skill format natively, so no special system prompt or “skill-writing” tool is required. The model extracts the reusable pattern from the interaction and structures it correctly.

One important detail: the description comes last. Only after the workflow is complete, the real trigger phrases become clear: terms like “deck”, “slides”, “presentation for a client”, or “pitch”. This ordering, workflow first, description second, is a key departure from traditional documentation practices.

A final review pass is still useful. Claude can check whether SKILL.md includes explanations it already knows, such as Python import syntax, what a .pptx file is, or general Drive authentication behavior. Those details should be removed.

Ideally, SKILL.md remains under roughly 500 lines and reads like a table of contents. It should point to supporting files e.g. brand-guidelines.md for visual constraints, google-slides-setup.md for configuration, and the scripts directory for execution, and stop there. Supporting files are loaded on demand, so they add no overhead until they are needed.

3.2 Deploying the Skill

There are four surfaces, each with a completely separate skill system. They do not sync with one another.

• Claude.ai and the Claude desktop app share the same installation method: zip the skill directory with SKILL.md at the root and upload it via Settings → Features → Skills. Skills installed this way are personal, no org-wide distribution exists, so each team member uploads their own copy. It requires Pro, Max, Team, or Enterprise with code execution enabled. More information in the documentation: Using Skills in Claude

• Claude Code (the command-line tool) uses the local filesystem instead. Place the skill at ~/.claude/skills/branded-slides/ and Claude Code discovers it automatically at startup. This is the right place to develop and iterate. When stable, move it to .claude/skills/branded-slides/ inside a project repository and commit it: every teammate who pulls the repo gets the skill immediately, no setup required. More information in the documentation: Agent Skills in Claude Code

• The Claude API is for automated pipelines. Skills are uploaded via /v1/skills endpoints and become workspace-wide. The important constraint here is that the API container has no outbound network access, which means for instance, in our use case, that the Drive upload step will not work as-is. More information in the documentation: Using Skills with the Claude API

3.3. Iterating With Claude as Co-Author

A skill is rarely correct on the first attempt. The recommended approach uses two Claude contexts: one for authoring, one for testing.

In the authoring context, the current SKILL.md is shared and the observed issue is described. For example: “Claude skipped the executive summary slide.” Claude analyzes the behavior and proposes a fix. The change is then validated in a fresh testing session with the updated skill loaded. This loop, observe, revise, retest, continues until the behavior stabilizes.

Most issues fall into a few predictable categories:

• The description misses trigger phrases users naturally use

• An important constraint is buried too deep in SKILL.md

• A reference file is structured in a way that Claude reads only partially

Once identified, each of these problems is usually straightforward to fix.

Changes can be tracked directly in SKILL.md using a simple changelog block, which keeps the iteration history visible without requiring additional tooling.

Key Takeaways

✓ Prompt engineering does not scale as a workflow foundation: Large system prompts are difficult to version, fragile to small wording changes, and hard to maintain across teams. As workflows grow more complex, prompt-based control becomes increasingly brittle.

✓ Skills package expertise, not instructions: A skill is a structured folder containing a workflow the model can discover and load on demand. The shift is conceptual: instead of telling the model what to do each time, you document the procedure once and let the model apply it whenever relevant.

✓ Progressive loading keeps skills lightweight: The model loads only what it needs, when it needs it. Instructions, brand guidelines, and scripts stay dormant until triggered, adding no token overhead to unrelated tasks.

✓ Build the workflow first, write the skill second: The most reliable path to a good skill is running the task in Claude first, then extracting the reusable pattern. Trigger phrases, constraints, and edge cases only become clear once the workflow actually works.

✓ Skill development is an iterative process: Effective skills emerge through a loop of observing model behavior, refining SKILL.md, and testing again in a clean session until the workflow stabilizes.

✓ Skills turn LLM workflows into maintainable infrastructure: Because they live in the filesystem or a repository, skills can be versioned, shared across projects, and maintained alongside the rest of the codebase.

References

[1] Accelirate. Getting Started with Prompt Engineering: A Technical Guide for Developers. 2024. https://www.accelirate.com/prompt-engineering-guide-for-developers

[2] PromptLayer Blog. Disadvantage of Long Prompt for LLM. September 2025. https://blog.promptlayer.com/disadvantage-of-long-prompt-for-llm

[3] DigitalOcean. ChatGPT vs Gemini: How AI Assistants Stack Up in 2026. https://www.digitalocean.com/resources/articles/gemini-vs-chatgpt

[4] Zhang, B., Lazuka, K., Murag, M. (Anthropic). Equipping agents for the real world with Agent Skills. October 2025. https://www.anthropic.com/engineering/equipping-agents-for-the-real-world-with-agent-skills

[5] Anthropic. Agent Skills — Claude Code Docs. https://code.claude.com/docs/en/skills

[6] Jackson, F. Custom GPTs vs. Gemini Gems: A Technical Comparison. SSRN, December 2025. https://papers.ssrn.com/sol3/papers.cfm?abstract_id=5903642

This article is also available as a podcast! If you’re on the go or just want to absorb the content in audio format, you can listen to the full episode below 👇 The podcast is also available on Spotify and Apple Podcasts.

A growing blind spot is emerging in website analytics: the user intent behind AI assistant traffic.

Every day, ChatGPT, Claude, Perplexity, and other AI tools visit documentation, scan product pages, and access content on behalf of real users searching for answers. While modern analytics platforms have sophisticated bot filtering mechanisms that automatically exclude spam bots and malicious crawlers [1], they face a different challenge with AI assistants: these tools represent genuine user needs, yet their traffic patterns are fundamentally different from direct human browsing.

The problem isn’t simply about filtering, it’s about context loss. When a user asks an AI assistant “How do I integrate payment processing with [product X]?” or prompts it to “Compare pricing plans for [service Y],” the resulting bot traffic to those websites represents authentic user intent. And with assistants increasingly grounding their answers in real-time web search and retrieval, not just relying on their pre-trained LLM weights, these visits are happening continuously, invisibly, and with growing frequency.

Moreover, most users receive the final answer directly inside the AI interface without ever clicking through to the original site [2]. This creates an “invisible visit” phenomenon where content provides value, AI systems access and process the information, but traditional analytics either record nothing or strip away the critical context of what the user was actually trying to accomplish.

Recent research suggests AI crawler traffic accounts for 5-10% of total server requests for knowledge-based websites [2], yet this traffic is either aggregated into generic “bot” categories or excluded entirely. The distinction matters: newer AI crawlers often aren’t covered by standard bot filtering lists [3], and even when detected, analytical frameworks treat them like search engine crawlers rather than proxies for human questions and needs.

Objective

This article presents a practical framework for extracting actionable intelligence from AI bot traffic, treating it not as noise to filter out, but as a complementary signal that reveals user intent.

After reading this article, you will understand:

1. Why does AI bot traffic contain different intelligence than traditional analytics? How AI-mediated user interactions fundamentally differ from direct human browsing, what makes them valuable despite being non-human traffic, and why current analytics frameworks fail to capture this value

2. What insights can be extracted from AI bot access patterns? Which specific patterns in bot traffic reveal user questions, documentation gaps, and product understanding issues, supported by analysis of real-world data

3. How can organizations operationalize AI bot insights? Practical methods to segment, analyze, and integrate AI bot data into product, content, and marketing workflows, with concrete examples from actual implementation

Whether you’re a product manager trying to understand user needs, a data analyst seeking new signal sources, or a content strategist optimizing documentation, this framework will help you see AI bot traffic not as a measurement problem, but as an untapped source of user intelligence.

This article is part of a series exploring the evolving landscape of AI-powered search and content discovery, available at https://siteline.ai/blog/. SiteLine is a platform that tracks and optimizes product and service visibility across the AI search ecosystem, providing comprehensive insights into how brands perform in AI-powered responses from ChatGPT, Perplexity, Gemini, Claude, and other AI assistants.

Prerequisites: No prior knowledge is required. The concepts are explained in and accessible way for anyone interested in how AI assistant traffic can inform business and content strategy.

Subscribe now

1. The Framework in Brief

The Core Innovation: Why LLMs Beat Clustering

Most teams try to understand user behavior with URL rules, regex, or embedding-based clustering. The problem: these methods are descriptive, not interpretive. They group hits that look similar on the surface, but they can’t reconstruct the story behind a session.

For instance, basic clustering might lump together a “Pricing” page visit and a “Login” attempt simply because they share common keywords or happened in sequence. This means you can’t distinguish between a new prospect moving through your sales funnel and an existing customer performing routine tasks. The critical difference between someone evaluating your product versus someone already using it is completely invisible in your analytics.

This framework flips the approach:

• Instead of pre-defined rules, we use an LLM as a reasoning engine.

• Instead of isolated pageviews, we analyze full journeys: which pages were visited, in what order, and for how long.

• Instead of a fixed taxonomy, we let the model create dynamic tags and categories as it discovers new patterns in the traffic.

• Instead of generic tags, we ask the model to interpret the session: who is behind this query, and what are they trying to do?

In other words, we move from “this session visited /pricing and /docs” to “this is a developer evaluating integration feasibility based on API limits, then checking whether the pricing fits their use case.”

How It Works: The 4-Stage Process

The framework treats every AI bot visit as a breadcrumb trail, reconstructing them into coherent sessions to decode the human intent behind the prompt.

Stage 1: Separating Signal from Noise The analysis starts at the source: server logs. Unlike traditional web analytics that rely on JavaScript tracking pixels and client-side sessions, AI assistant visits leave traces in the raw HTTP request logs maintained by your hosting infrastructure, whether that’s an edge service like Vercel, a CDN like Cloudflare, or your origin servers.

A typical log entry might look like this:

185.72.144.53 yourwebsite.com - [12/Feb/2025:14:32:09 +0000]
"GET /features/our-brand-new-awesome-feature HTTP/1.0" 200 64312 "-"
"Mozilla/5.0 (Linux; Android 6.0.1; Nexus 5X Build/MMB29P) AppleWebKit/537.36
(KHTML, like Gecko) Chrome/97.0.4692.71 Mobile Safari/537.36
(compatible; Googlebot/2.1; +http://www.google.com/bot.html)"

Raw data is filtered to keep only AI assistants acting on real user queries, removing noise from SEO crawlers, uptime monitors, spam bots, and internal automation.

Stage 2: Reconstructing User Journeys > Individual hits are grouped by IP address and temporal proximity. A natural break in activity (e.g., a 5-minute gap) signals the start of a new session. This transforms your bot logs from a list of disconnected URLs into a collection of browsing narratives, each representing someone’s attempt to understand something about your product or service.

Stage 3: Semantic Enrichment > To decode cryptic URLs, the system scrapes every unique page visited and uses a lightweight LLM to generate a semantic title and summary. This translates technical paths (e.g., /docs/v2/auth) into a clear sequence of consumed concepts, creating a human-readable storyline for the LLM to analyze.

Stage 4: Inferring User Intent > A reasoning LLM analyzes the full enriched journey to answer three core questions, assigning a confidence score to each classification:

1. Topic: What is the precise subject? (e.g., “Payment Integration Errors”).

2. Role: Who is the user? (e.g., Developer vs. Decision-Maker).

3. Goal: What are they trying to do? (e.g., Evaluation vs. Active Implementation).

This final stage converts raw data into dynamic, quantified insights that evolve automatically as new content and user behaviors emerge.

2. Case Study: What the Data Actually Reveals

Let’s see what happens when the framework is applied in practice. For this analysis, server logs were collected from three pilot websites over a 24-hour period and filtered to isolate only ChatGPT bot traffic. For this analysis, we’ll go in depth with Tradevision.io, a data platform for stock and options trading.

Anatomy of an AI-Mediated Session

Before examining aggregate patterns across thousands of sessions, let’s trace a single journey to understand what session reconstruction and LLM-based classification reveals and why it matters.

The framework’s LLM classifier analyzed this sequence and produced:

- Topic: Trading platform review
- User Role: decision-maker (confidence: 90/100)
- Intent Type: comparison (confidence: 92/100)
- Reasoning: The session comprises multiple TradeVision blog pages focused on comparing and evaluating trading platforms (AI-driven tools, 2025 platform rankings, mobile vs desktop, and platform reviews). This indicates a decision-making, comparison-oriented intent rather than learning or implementation.

This session likely reflects a user asking something like “What’s the best AI trading platform for mobile in 2025?” The assistant gathered comparative information step by step, moving from general guidance to specific recommendations to product details.

The next sections scale this analysis across thousands of sessions to reveal broader patterns in user segmentation, content performance, and optimization priorities.

Topic Distribution: Identifying Content Opportunities

The traffic data reveals two dominant themes: strong demand for options trading education and consistent interest in platform comparison content. Understanding these patterns helps identify where content investment delivers the highest return for tradevision.io in AI search visibility.

Figure 3 shows that options trading content captures over 400 sessions across three related topics: “Options trading strategies” (153), “Options trading toolkit” (144), and “Options trading concepts” (119). Together, these three topics account for 47% of all sessions, out of a total of 873.

The topic-intent relationship displayed in Figure 4 shows distinct patterns. Options content attracts pure learning intent while “Trading platform review” shows balanced distribution, indicating the content successfully serves both discovery and evaluation stages.

User Segmentation: Understanding User Roles, Intent, and Behavioral Patterns

The findings become even stronger when we analyze the user-level segmentation.

Figure [8] shows that end-users represent the majority of traffic and demonstrate 93.9% learning intent with minimal comparison (5.4%). Decision-makers show the inverse pattern: 60.4% comparison intent with secondary learning (39.6%).

These fundamentally different browsing behaviors reflect different stages in the consideration journey, end-users building knowledge, decision-makers actively evaluating alternatives.

3. Why This Matters: Strategic Implications

The Zero-Click Blind Spot

For Tradevision, AI assistants account for 17% of all sessions to content about options trading strategies. Traditional analytics, however, register zero visits, zero engagement, zero value. And yet, real users received answers, while AI systems began associating TradeVision with expertise in options trading.

This is the core challenge: content is delivering business value without activating conventional success metrics. Companies focused only on direct traffic are measuring a shrinking slice of how users actually discover, evaluate, and understand their products. This framework surfaces that hidden demand.

Early Signal Detection as Competitive Edge

AI assistants capture user intent long before a prospect commits to visiting a website, trying a product, or contacting support.

Consider the case of TradeVision. By analyzing AI search behavior, they detected a surge of decision-makers on platform comparison pages but there was near-zero intent to actually implement the software. Traditional funnel analysis would simply see this as “bounced traffic,” but AI signals revealed a potential issue: a hidden onboarding gap?

This kind of signal allows for a predictive strategy:

• Fix friction before it escalates into support tickets.

• Expand content before feature requests pile up.

• Claim ownership of topics before competitors even notice the trend.

The Compound Effect: Most importantly, this first-mover advantage compounds over time. When AI assistants consistently cite TradeVision for options trading, they build deep mental associations. Competitors entering the market later face a steeper challenge: they must not only match the product but also unseat an established algorithmic bias in your favor.

AI Discoverability: The New Game

While traditional SEO focuses on visibility on a search engine results page, AI discoverability focuses on retrieval and citation. The goal is to be the primary source used by an assistant to generate an answer, often without the user ever visiting a website.

This shift demands that we prioritize structure over keywords.

TradeVision exemplifies this success. Their content performs well not due to backlink volume, but because it is architected for synthesis as it provides self-contained, comprehensive answers. Conversely, content strategies that fragment information across multiple pages perform poorly in AI-mediated discovery.

The proposed framework allows companies to audit their digital footprint, identifying which assets successfully serve this new discovery layer and which are invisible to the algorithms that matter most.

What’s Next?

AI assistant traffic represents a paradigm shift in how we understand user behavior online. Traditional analytics tell us what users do when they visit our sites, which pages they view, how long they stay, and where they click. AI bot traffic, when properly analyzed, reveals something more fundamental: what users are trying to accomplish before they even decide whether to visit.

These insights can be operationalized across teams.

• Product managers can spot recurring workarounds or friction points that assistants repeatedly access and prioritize features accordingly.

• Content teams can identify high-traffic topics where assistants still struggle to find clear answers, signaling documentation gaps.

• Marketing and sales can understand the comparison patterns and evaluation criteria users rely on long before they land on the pricing page.

• Engineering teams can uncover technical blockers such as rate limits, authentication quirks, crawler-unfriendly endpoints [1] [5].

There are limitations. Session attribution remains imprecise, as a single user query can trigger multiple bot requests across different pages. Intent inference has natural boundaries: sequences reveal likely goals, but without the original prompt, certainty is impossible. And AI behavior itself evolves as models and retrieval systems improve, requiring the framework to adapt continuously.

References

[1] Google Analytics Help. “[GA4] Known bot-traffic exclusion.”

[2] Writesonic. (2025). “Introducing AI Traffic Analytics: Track Traffic from ChatGPT, Gemini & Perplexity.”

[3] IronPlane. (2025). “How to Clean Up Your Google Analytics Data from AI Bot Traffic.“

[4] Vlad Zotov, The AI Crawler Optimization Guide to Increase Site Visibility, Siteline

[5] Giacomo Zecchini, Alice Alexandra Moore, Malte Ubl, Ryan Siddle, The rise of the AI crawler, Vercel

This article is also available as a podcast! If you’re on the go or just want to absorb the content in audio format, you can listen to the full episode below 👇 The podcast is also available on Spotify and Apple Podcasts.

Most agent systems are effectively stateless: they reason well in the moment but forget everything once execution ends. As agents move toward longer-running workflows and repeated interactions, this limitation becomes a structural bottleneck. Long-term memory is what makes agents able to personalize, improve, and coordinate over time, moving them from reactive pipelines into adaptive systems.

Objective

This article explores how to incorporate long-term memory into LangGraph agent architectures. It discusses how durable memory can be designed, stored, and accessed within graph-based execution flows without sacrificing transparency or debuggability.

After reading this article, you will understand:

1. How long-term memory differs from short-term context: The distinction between ephemeral conversation history and persistent knowledge, and how different memory types map to LangGraph agent architectures.

2. Where LangGraph’s native memory stops: What checkpointers and the Store API provide for persistence—and their critical limitations for structured extraction, validation, and intelligent merging.

3. How to implement production-ready memory with Trustcall: Using practical examples, you’ll learn to build schema-driven memory systems that extract, merge, and persist user profiles and collections reliably across sessions.

This article is part of a broader series on building agent systems with LangGraph. Short-term memory is already covered in Beyond the Demo: Building AI Agents That Remember, Recover, and Scale (Part 2), which provides additional context on the topic.

Prerequisites: Basic knowledge of Python, LLMs, and prompt engineering is recommended. Readers should either be familiar with LangChain concepts (chains, models, prompts) or have read Part 1 of this series (or listen the podcast edition) to understand the foundations of LangGraph’s graph, node, and state model.

Tools & libraries: LangGraph, LangChain, LangSmith, OpenAI, arXiv API

You can find the code here on GitHub. The demo video is at the end of the article!

Subscribe now

1. Foundations of Long-Term Memory in Agent Systems

Long-term memory in agents is not a single component you “add” to an LLM. It is an architecture that separates what must be present right now (for reasoning) from what must persist over time (for personalization, continuity, and learning).

This section introduces the main memory types used in agent systems, how they interact during execution, and the practical constraints that shape real implementations.

1.1 The Four Core Memory Types in Agent Architectures

Most long-term memory designs for LLM agents can be explained with four complementary memory types, often discussed in cognitive-inspired frameworks like CoALA [1]. The key idea is simple: different information has different “lifecycles”, so it should not all be stored (or retrieved) the same way.

1. Working memory (context window): It’s the agent’s immediate awareness. The tokens currently inside the model’s context window. It holds what the agent is actively reasoning over right now, recent turns, task instructions, intermediate artifacts, and recently retrieved facts. Even with long-context models, this remains the most constrained and expensive resource, so agents need explicit strategies for what stays in-context and what gets evicted [2].

2. Semantic memory (facts and knowledge): It stores distilled knowledge: facts, concepts, preferences, without needing the full story of when they were learned. In agent systems, this is where many RAG-style approaches live: embeddings in vector databases, structured fact stores, or knowledge graphs.

3. Episodic memory (interaction history): It preserves sequences of events as they happened: full conversations, task trajectories, ordered observations. Unlike semantic memory, it keeps narrative context and temporal flow. This is valuable for audit trails, debugging, pattern analysis over conversations, and reflection loops that synthesize higher-level insights from experience [4]. In practice, episodic memory often looks like a message store and/or indexed conversation chunks, sometimes designed explicitly to store what was evicted from working memory [2].

4. Procedural memory (skills and behaviors): This represents the knowledge of how to do things: the agent’s skills, tool usage patterns, and stable workflows. For LLM agents, procedural memory is typically encoded in model weights, system prompts, and tool/function registries. It changes more slowly than other memory types and usually requires explicit updates (deployments, fine-tuning, tool additions). Recent work also explores agents that can learn reusable “skills” from experience rather than only storing facts [5].

Note: In production, many systems also need “prospective memory” to remember to execute intended future actions (follow-ups, reminders, scheduled steps). This is often implemented via task queues, reminders, and goal trackers rather than retrieval-based memory.

1.2 Coordination Between Memory Types

The usefulness of memory is not just storage but also coordination: how information moves between memory types during execution.

• Semantic → Working (retrieve before reasoning): Before an agent responds or acts, it typically retrieves relevant facts from semantic memory and injects them into working memory. This is the core RAG pattern: keep long-lived information outside the context window, then pull only what’s needed for the current decision.

• Working → Episodic (persist what happened): As interactions unfold, the agent should persist the event sequence to episodic storage. This creates a reliable record for later debugging, audits, or trajectory-based learning. OS-inspired memory systems explicitly treat the context window as scarce “RAM” and page information in/out across memory tiers [2].

• Episodic → Semantic (consolidate experiences into facts): Over time, raw experience becomes more useful when summarized into stable knowledge. Generative Agents [4] popularized “reflection” mechanisms that periodically synthesize episodic memories into higher-level insights, which can then be stored as semantic memory and reused later. Mem0 formalizes a similar idea as a production-oriented pipeline: extract and consolidate important information from conversations into durable memories [3].

• Procedural + Semantic (apply skills to knowledge): Procedural memory provides the capability; semantic memory provides the content. For example, a coding agent uses procedural knowledge (how to refactor, how to call tools) to apply semantic knowledge (project requirements, architecture constraints).

• Time as a shared dimension: A subtle but important point is that “truth” changes over time: preferences drift, jobs change, policies update. Temporal memory layers explicitly track when facts were true and when they became invalid, enabling queries like “what was true then?” vs “what is true now?” [6]. This is especially important when semantic memories are mutable rather than append-only.

1.3 Practical Limits & Trade-offs of Long-Term Memory in Agent Architectures

Once memory is introduced, agent design becomes a balancing act across cost, correctness, and control.

• Limited working memory does not disappear with long context:

  • Longer context windows reduce friction, but they do not remove the need for memory management.

  • Models can struggle when relevant information is buried deep inside large contexts (“lost in the middle”), and performance can degrade based on where the key information appears [7].

  • Practically, large contexts also increase latency and cost, so it is rarely optimal to treat the context window as a full-history database.

• Storing everything as “semantic memory” increases noise

  • A common failure mode is embedding all messages, logs, and metadata into a single vector index. This often creates retrieval noise, loses narrative structure, and becomes expensive to search at scale.

  • Production approaches therefore separate concerns: keep full transcripts and sequences as episodic memory, and store extracted, consolidated facts as semantic memory [3].

• Long-term memory must handle change, not just accumulation

  • Without explicit invalidation, agents surface stale or contradictory information. Temporal modeling addresses this by attaching validity semantics to memories (e.g., when a fact became true vs when it stopped being true), which helps prevent outdated facts from dominating retrieval [6]. This is one reason temporal knowledge-graph approaches have become more common for agent memory layers.

• Retrieval quality is often the bottleneck

  • If retrieval brings the wrong memory into working context, downstream reasoning will amplify it. This is why many systems combine techniques (metadata filtering, hybrid lexical + dense search, reranking, structured retrieval) rather than relying on a single similarity search mechanism. Benchmarks and production architectures increasingly focus on measuring memory quality and reducing retrieval-driven errors [3].

• Efficiency matters at production scale

  • Memory has an operational footprint: storage, embedding, indexing, reranking, and retrieval all accumulate cost. Architectures like Mem0 report substantial reductions in latency and token usage compared to full-context approaches, by storing distilled facts and retrieving selectively [3]. OS-inspired designs similarly motivate tiered memory and paging strategies to keep working memory small and relevant [2].

Building production-grade long-term memory is often more challenging than building the agent itself, and the effort required to make memory reliable, scalable, and correct is frequently underestimated. The next section shows how LangGraph supports the practical implementation of long-term memory in agent systems.

2. Implementing Long-Term Memory with LangGraph

2.1 Use Case: A Research Paper Recommender with Persistent User Memory

Consider a research paper recommendation system that helps users discover recent publications aligned with their interests. Unlike a simple one-shot query system, this agent needs to remember user preferences across multiple interactions to provide increasingly personalized recommendations.

The Challenge: The agent must maintain two types of long-term memory:

1. Profile-based memory: Track evolving user interests, expertise level, and reading history

2. Collection-based memory: Organize recommended papers into topic-based collections that grow over time

For example, when a user first queries “show me papers on transformers,” the agent should:

• Extract and store “transformers” as an interest

• Fetch relevant recent papers from ArXiv

• Create a collection grouping papers by subtopic

In subsequent interactions, when the user asks “what about quantum computing papers?”, the agent should:

• Add “quantum computing” to their existing interests (not replace them)

• Remember their expertise level for appropriate paper selection

• Maintain separate collections for both topics

This requires persistent, structured memory that can be incrementally updated without losing previous information, a capability that goes beyond simple conversation history.

You can find the code here on GitHub.

2.2. Built-in Memory Options in LangGraph & Their Limitations

LangGraph provides memory mechanisms at different levels.

For short-term memory, which consists of the conversation context within a single session, LangGraph checkpointers and message history work well. These were covered in detail in my previous article on LangGraph fundamentals: Beyond the Demo: Building AI Agents That Remember, Recover, and Scale (Part 2)

However, for long-term memory that persists across sessions and user interactions, LangGraph introduced the Store API [8].

The Store API provides a dedicated interface for managing persistent, cross-thread memory [9]:

from langgraph.store.memory import InMemoryStore

# Create a store for long-term memory
store = InMemoryStore()

# user_id would come from your application context (e.g., session, authentication)
user_id = "user_123"  # Example: unique identifier for the current user

# Store user profile
store.put(
    namespace=("user_profiles", user_id),
    key="profile",
    value={"interests": ["transformers"], "expertise_level": "beginner"}
)

# Retrieve later (in a different session/conversation)
profile = store.get(namespace=("user_profiles", user_id), key="profile")

Key Benefits:

The Store API offers several advantages [9]:

• Persistence: Data survives beyond individual conversation threads

• Namespacing: Organize memory by user, topic, or custom categories

• Cross-thread access: Share memory across different conversation sessions

Limitations of the Store API

While the Store API handles persistence, it has significant limitations for complex memory management:

• No structured extraction: Store is a key-value database—it doesn’t extract structured information from conversations. You must manually parse and structure data before storing [2][3]

• No intelligent merging: You’re responsible for all merge logic, deduplication, and conflict resolution when updating existing memory [9]

• No schema validation: Store accepts any data structure with no type checking or field validation, risking schema drift over time [8]

• Manual extraction pipeline: Each memory update requires custom code for prompting, LLM calls, parsing, validation, retrieval, merging, and storage [10]

The Core Problem: LangGraph’s Store API solves persistence but not structured memory management. This is where Trustcall bridges the gap.

2.3. Implementing Long-Term Memory with Trustcall

Trustcall provides reliable structured extraction for LangGraph agents through its create_extractor function [11]. It handles schema validation, duplicate detection, and incremental updates. This eventually transforms the complex task of memory management into a few lines of code.

Profile-Based Memory: Tracking User Preferences

Profile-based memory maintains a structured representation of user attributes that evolves over time. Here’s how to implement it with Trustcall:

Step 1: Define the Memory Schema

from pydantic import BaseModel, Field

# 1. Define schema for user profile
class ResearchProfile(BaseModel):
    """User's research profile"""
    interests: List[str] = Field(description="Research topics of interest")
    expertise_level: str = Field(description="beginner, intermediate, or expert")
    read_papers: List[str] = Field(default=[], description="Papers already read")

Step 2: Create the Trustcall Extractor

from trustcall import create_extractor
from langchain_openai import ChatOpenAI

# 2. Create extractor
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

profile_extractor = create_extractor(
    llm,
    tools=[ResearchProfile],
    tool_choice="ResearchProfile"
)

The create_extractor function creates a specialized agent that reliably extracts structured data matching your Pydantic schema [5]. By setting tool_choice, you ensure the extractor always attempts to populate your UserProfile schema.

Step 3: Extract profile

# 3. Extract initial profile

conversation = """
I'm interested in transformers and attention mechanisms.
I'm a beginner in this field.
"""

result = profile_extractor.invoke({
    "messages": [{"role": "user", "content": conversation}]
})

profile = result["responses"][0]
print("Initial profile:")
print(profile.model_dump_json(indent=2))

Output

Initial profile:
{
  "interests": [
    "transformers",
    "attention mechanisms"
  ],
  "expertise_level": "beginner",
  "read_papers": []
}

# 4. Update profile with new information
new_conversation = """
I've been reading about quantum computing lately.
I think I'm getting better at this, maybe intermediate level now.
"""

updated_result = profile_extractor.invoke({
    "messages": [{"role": "user", "content": new_conversation}],
    "existing": {"ResearchProfile": profile.model_dump()}
})

updated_profile = updated_result["responses"][0]
print("\\nUpdated profile:")
print(updated_profile.model_dump_json(indent=2))

Output

Updated profile:
{
  "interests": [
    "transformers",
    "attention mechanisms",
    "quantum computing"
  ],
  "expertise_level": "intermediate",
  "read_papers": []
}

Key Benefits:

• Automatic validation: Trustcall ensures extracted data matches your Pydantic schema [12]

• Reliable parsing: No manual JSON parsing or error handling needed

• Merge control: You control merge logic (union, replace, append) while Trustcall handles extraction

Collection-Based Memory: Growing Lists of Structured Data

Collection-based memory maintains growing lists where new items are added over time. Trustcall’s enable_inserts=True feature makes this seamless [13]:

Step 1: Define the Collection Schema

# 1. Define schema for paper notes
class PaperNote(BaseModel):
    """A note about a research paper"""
    paper_title: str = Field(description="Paper title")
    key_insight: str = Field(description="Main takeaway from the paper")
    relevance: str = Field(description="Why this paper matters")

Step 2: Create Extractor with Insert Support

# 2. Create extractor with enable_inserts=True
collection_extractor = create_extractor(
    llm,
    tools=[PaperNote],
    tool_choice="PaperNote",
    enable_inserts=True  # KEY: allows adding new items
)

The enable_inserts parameter allows the extractor to create new instances of your schema while preserving existing ones [13]. This is essential for maintaining growing collections.

Step 3: Manage Collections with Existing Data

# 4. Add new papers to collection
conversation_with_new_papers = """
I just read two papers:

1. "Chain-of-Thought Prompting" - Shows that step-by-step reasoning 
   improves LLM accuracy on complex tasks.

2. "RAG: Retrieval-Augmented Generation" - Combines retrieval with 
   generation for better factual accuracy.
"""

result = collection_extractor.invoke({
    "messages": [{"role": "user", "content": conversation_with_new_papers}],
    "existing": existing_data
})

print("\\nUpdated collection:")
for response, metadata in zip(result["responses"], result["response_metadata"]):
    is_new = metadata.get("json_doc_id") is None
    status = "NEW" if is_new else "EXISTING"
    doc_id = metadata.get("json_doc_id", "N/A")
    
    print(f"\\n[{status}] {response.paper_title}")
    print(f"  Insight: {response.key_insight}")
    print(f"  Doc ID: {doc_id}")

Updated collection:

[NEW] Chain-of-Thought Prompting
  Insight: Step-by-step reasoning improves LLM accuracy on complex tasks.
  Doc ID: N/A

[NEW] RAG: Retrieval-Augmented Generation
  Insight: Combines retrieval with generation for better factual accuracy.
  Doc ID: N/A

Key Benefits:

• Incremental updates: enable_inserts=True allows adding new collections while preserving existing ones [13]

• Context awareness: The existing parameter tells Trustcall about current state [14]

• Intelligent grouping: The LLM can merge new papers into existing topics or create new ones

Putting It All Together

In a complete LangGraph agent, these memory extraction patterns integrate seamlessly into your workflow nodes. The research paper recommender demonstrates this integration across three sequential stages:

1. Profile Management Node: Uses the profile extractor to update user interests and expertise level from each query. The extracted profile is stored in the graph state and persists across conversation turns through LangGraph’s checkpointer [9].

2. Recommendation Node: Leverages the updated profile to fetch relevant papers from ArXiv. This node reads the current interests from state to determine search topics, demonstrating how extracted memory directly influences agent behavior.

3. Collection Management Node: Uses the collection extractor with enable_inserts=True to organize recommended papers by topic. By passing existing collections via the existing parameter, new papers are intelligently grouped with previous ones, building a growing knowledge base over time.

The key architectural insight is that Trustcall handles the extraction complexity within each node, while LangGraph orchestrates the flow between memory updates, tool calls, and user responses. The graph state serves as the persistence layer, with Trustcall ensuring that all updates to user_profile and collections are structured, validated, and properly merged.

For the complete implementation showing how these extractors integrate into LangGraph nodes and state management, see the full code example https://github.com/linafaik08/scientific-graph-agent.

Trustcall eliminates the boilerplate of structured memory management, letting you focus on your agent’s logic rather than extraction plumbing. The combination of Pydantic schemas for validation, intelligent extraction, and merge support makes long-term memory practical and reliable [11][12].

Key Takeaways

✓ Long-term memory is architectural, not prompt engineering: Reliable agent memory requires explicit design around what to store, where it persists, and how it updates. Larger context windows or appending conversation history don’t solve personalization or

scalability.

✓ Different memory types must be coordinated: Working, semantic, episodic, and procedural memory each have distinct lifecycles. Effective agents emerge from controlled flows between them: episodic consolidation into semantic memory, semantic retrieval into working memory.

✓ LangGraph makes memory explicit but not automatic: State and checkpointing provide transparency and debuggability, but don’t solve structured extraction, schema enforcement, or conflict resolution. These must be handled explicitly in your nodes.

✓ Schema-driven updates prevent memory drift: Typed schemas (Pydantic models) for profiles and collections enable controlled updates, deduplication, and validation. This prevents memory from becoming noisy append-only logs.

✓ Production memory is harder than agent logic: Most agent failures stem from incorrect retrieval, stale facts, or uncontrolled memory growth, not reasoning quality. Building accurate, scalable memory typically demands more effort than the core agent logic itself.

References

1. CoALA — Cognitive Architectures for Language Agents (arXiv:2309.02427) https://arxiv.org/abs/2309.02427

2. MemGPT — Towards LLMs as Operating Systems (arXiv:2310.08560) https://arxiv.org/abs/2310.08560

3. Mem0 — Building Production-Ready AI Agents with Scalable Long-Term Memory (arXiv:2504.19413) https://arxiv.org/abs/2504.19413

4. Generative Agents: Interactive Simulacra of Human Behavior (arXiv:2304.03442) https://arxiv.org/abs/2304.03442

5. Letta — Skill Learning: Bringing Continual Learning to CLI Agents https://www.letta.com/blog/skill-learning

6. Zep — A Temporal Knowledge Graph Architecture for Agent Memory (arXiv:2501.13956) https://arxiv.org/abs/2501.13956

7. Lost in the Middle: How Language Models Use Long Contexts (arXiv:2307.03172) https://arxiv.org/abs/2307.03172

8. LangGraph Documentation — Store API Overview https://langchain-ai.github.io/langgraph/concepts/memory/#store

9. LangGraph Documentation — Using the Store https://langchain-ai.github.io/langgraph/how-tos/manage-cross-thread-memory/

10. The AI Practitioner — Beyond the Demo: Building AI Agents That Remember, Recover, and Scale (Part 2), Lina Faik

11. Trustcall Documentation — create_extractor https://github.com/Sizaar/trustcall

12. Trustcall Documentation — Schema Validation https://github.com/Sizaar/trustcall#schema-validation

13. Trustcall Documentation — enable_inserts https://github.com/Sizaar/trustcall#incremental-updates

14. Trustcall Documentation — Passing Existing Data https://github.com/Sizaar/trustcall#managing-existing-data

This article is also available as a podcast! If you’re on the go or just want to absorb the content in audio format, you can listen to the full episode below 👇 The podcast is also available on Spotify and Apple Podcasts.

As agent systems grow from simple chains to complex multi-agent setups, the graphs meant to bring clarity end up creating confusion. What begins as a clean five-node workflow explodes into 30+ nodes, each relying on unclear state changes.

The key to scaling graph complexity is understanding when to parallelize independent work, when to encapsulate logic in subgraphs, and when to let the system dynamically spawn tasks based on runtime conditions.

Objective

This article focuses on solving three scalability problems: latency from sequential execution, complexity from oversized graphs, and rigidity from static workflows.

After reading this article, you will understand:

1. How to reduce latency through parallelization while managing concurrent state updates?

2. How to structure multi-agent architectures using subgraphs with isolated or shared state?

3. How to implement map-reduce patterns for variable workloads using the Send API?

This article is the fourth part of the LangGraph series on graph-based agent architectures. While earlier parts focused on individual graph components, Part 4 explores how to compose and orchestrate multiple agents efficiently.

Prerequisites: Basic knowledge of Python, LLMs, and prompt engineering is recommended. Readers should either be familiar with LangChain concepts (chains, models, prompts) or have read Part 1 of this series (or listen the podcast edition) to understand the foundations of LangGraph’s graph, node, and state model.

Tools & libraries: LangGraph, LangChain, LangSmith, OpenAI, arXiv API

You can find the code here on GitHub.

Subscribe now

1. Understanding LangGraph’s Concurrency Models

1.1 When should parallelization be chosen over sequential execution?

What is parallelization in LangGraph?

Parallelization in LangGraph means executing multiple nodes simultaneously rather than one after another. Instead of waiting for one task to complete before starting the next, multiple tasks launch at the same time and run concurrently.

This is particularly valuable when making multiple independent API calls, database queries, or LLM requests that don’t depend on each other’s results.

In LangGraph, parallel execution is created by adding multiple edges from a single node to multiple destination nodes. The graph automatically detects this fan-out pattern and executes the destination nodes concurrently in what LangGraph calls a “superstep.”

Sequential vs. Parallel: A concrete example

Consider building a research assistant that answers questions by consulting both ArXiv (for academic papers) and Wikipedia (for general knowledge).

Sequential execution (one after another):

With sequential execution, if ArXiv takes 2 seconds and Wikipedia takes 2 seconds, the total wait time is 4 seconds.

Parallel execution (both at once):

With parallel execution, both APIs are called simultaneously. The total wait time becomes only 2 seconds (the duration of the longest call), cutting execution time in half.

Comparisons:

In the research paper use case, running both tools sequentially took 61.46s, compared to just 0.45s when executed in parallel, a 137× speedup!

Key implementation considerations

• The fundamental requirement for parallelization is independence. The ArXiv and Wikipedia calls work in parallel because neither needs the other’s output to proceed. When nodes truly run independently, parallelization delivers immediate performance gains.

• However, state management requires careful attention. If both nodes update the same state key (like adding results to a “documents” list), reducers become essential to merge their updates properly. As discussed in the previous blog post on reducers (cf. Part 2 of the series), these functions define how concurrent updates combine.

• Understanding supersteps is crucial for managing parallel execution. A superstep is an execution unit in LangGraph that groups together nodes that can run concurrently. In our previous example, the graph detects multiple edges fanning out from a single node as ArXiv and Wikipedia starting simultaneously. It creates a superstep containing those parallel nodes. All nodes in a superstep execute at the same time and must all complete before the graph proceeds to the next step.

• Error handling requires special consideration within supersteps. If one parallel node fails, the entire superstep fails atomically. This means if Wikipedia succeeds but ArXiv fails, neither result gets saved to state. This transactional behavior prevents inconsistent states but requires robust retry logic or fallback strategies in the node design. When using checkpointing, however, LangGraph saves results from successful nodes internally, so only failing branches need to retry when the graph resumes.

1.2 What are the performance trade-offs of concurrent node execution?

Speed gains vs. resource consumption

• Parallel execution reduces latency but increases instantaneous resource usage. In the ArXiv and Wikipedia example, wait time drops from 4 seconds to 2 seconds, but two API calls now happen simultaneously instead of sequentially. This shift creates several practical implications.

• Rate limits become a primary concern. If the Wikipedia API allows 60 requests per minute and five parallel searches run simultaneously, the quota gets consumed in 12 seconds instead of 1 minute. Monitoring and potentially throttling concurrent requests becomes necessary. LangGraph provides a max_concurrency configuration parameter to limit how many nodes can run at once, offering direct control over this trade-off.

Complexity vs. debuggability

• Parallel execution trades simplicity for speed. Sequential execution offers a clear linear path to trace when problems occur. Parallel execution requires tracking multiple simultaneous paths, understanding how states merge, and diagnosing which branch caused failures. If a synthesis node fails, determining whether bad data came from ArXiv, Wikipedia, or both requires investigation.

• Better logging and observability become essential. Execution traces show nodes completing in unpredictable order, making bug reproduction challenging. The trade-off is clear: parallelization delivers speed at the cost of operational complexity, making sense primarily when waiting on multiple slow external services that can run independently.

2. Subgraphs for Modular Agent Architectures

2.1 When to use isolated state vs. shared state between agents?

What are subgraphs in LangGraph?

Subgraphs are complete, self-contained graphs that can be embedded as nodes within a parent graph. They are reusable components or specialized agents that encapsulate their own logic and workflow. A subgraph is compiled independently and then added to a parent graph just like any other node, but internally it can have its own multi-step workflow, state schema, and execution logic. This architecture becomes particularly powerful for multi-agent systems. Instead of building one massive graph with dozens of nodes, the system can be decomposed into specialized agents (subgraphs), each responsible for a specific domain.

For example, a research assistant might have an orchestrator agent (parent graph) that delegates to specialized subgraphs: a literature review agent that searches and summarizes academic papers, a data extraction agent that pulls statistics and figures from documents, and a citation management agent that formats references. Each subgraph operates as an independent agent with its own internal workflow.

Shared state: the simple approach

The simplest subgraph pattern uses shared state keys between the parent and child graphs. According to the LangGraph subgraphs documentation, when a subgraph shares state keys with the parent graph, communication happens automatically through these shared channels. This is common in multi-agent systems where agents communicate over a shared “messages” key.

Consider a research system where an orchestrator agent (parent) delegates to specialized research agents (subgraphs):

With shared state, the parent’s messages flow directly into the subgraph, the subgraph’s internal nodes can add more documents, and these updates flow back to the parent. This creates seamless communication but also means the subgraph can see and modify all parent state, which may not always be desired.

Isolated state: maintaining boundaries

Isolated state becomes necessary when agents need private information that shouldn’t leak to other parts of the system. With isolated state, the subgraph has its own schema completely independent from the parent. This requires explicit state transformation at the boundaries:

The trade-off is clear: shared state offers simplicity and automatic communication, while isolated state provides encapsulation and privacy at the cost of manual state transformation. The choice depends on whether agent boundaries need enforcement or whether seamless information flow is more valuable.

2.2 How to maintain separate memory for each agent in a multi-agent system?

Independent checkpointing for agent memory

Each subgraph can maintain its own memory by compiling it with a dedicated checkpointer. This is particularly useful in multi-agent systems where each agent needs to track its own internal history independently.

This architecture enables agents to build up expertise and context over time as each agent maintains focused, relevant memory for its domain.

Memory visibility and debugging

The documentation notes that subgraph state can be viewed using graph.get_state(config, subgraphs=True), but with an important caveat: subgraph state is only accessible when the subgraph is interrupted. Once execution resumes, that checkpoint moves into the past and is no longer directly accessible.

This has practical implications for debugging and monitoring. To understand what happened inside a subgraph, interruption points need strategic placement, or the subgraph needs to explicitly surface key information back to the parent state. For example, a literature review agent might maintain detailed paper evaluations in its private state but surface only the final summary to the parent.

3. Map-Reduce with the Send API

3.1 When should dynamic branching be used vs. static parallelization?

What is map-reduce in LangGraph?

Map-reduce is a pattern for distributing work across multiple parallel tasks and then aggregating their results.

• The “map” phase splits work into independent units that process in parallel

• The “reduce” phase combines their outputs.

In LangGraph, map-reduce is implemented using the Send API, which enables dynamic task creation at runtime where the number and configuration of parallel tasks are determined by the graph’s state rather than fixed at design time.

Consider a research assistant with a query clarifier that interprets ambiguous research questions in multiple ways. The number of interpretations isn’t known until the clarification completes, and each interpretation needs separate processing.

The route_to_reaserchers function returns a list of Send objects, each specifying which node to invoke and with what state. LangGraph executes all these search tasks in parallel within a single superstep, then proceeds to synthesis once all complete.

The trade-off: flexibility vs. simplicity

• Map-reduce provides flexibility for variable workloads but introduces complexity. The graph structure isn’t fully visible at design time since the number of parallel tasks depends on runtime state. A query might generate 2 interpretations or 10, creating different execution patterns each run. This makes visualization and debugging more challenging.

• Static parallelization offers simplicity and predictability but lacks adaptability. The choice depends on whether the parallel structure is truly fixed. For research tasks where interpretation counts vary (sometimes needing 2 searches, sometimes 8), map-reduce is essential. For fixed API queries with consistent structure (always ArXiv + PubMed + Semantic Scholar), static parallelization is cleaner and more maintainable.

3.2 What happens when parallel branches have different execution lengths?

The challenge of uneven task durations

In static parallelization with fixed APIs, branches typically complete around the same time since they perform similar operations. With map-reduce, different tasks may require drastically different processing times. A research assistant searching multiple query interpretations might get instant results for one interpretation (cached or simple query) but wait 20 seconds for another (complex query requiring extensive processing).

The problem emerges when these tasks of different lengths need to converge. Consider a graph that searches multiple interpretations in parallel, then synthesizes findings. If the synthesis node starts as soon as the first search finishes, it works with incomplete data. The graph needs a way to ensure all parallel work completes before proceeding.

The defer parameter for synchronization

In LangGraph, when multiple parallel branches converge to the same node, the system automatically waits for all branches to complete before proceeding Langchain. This built-in synchronization is fundamental to how LangGraph’s superstep execution model works.

The synthesize node automatically waits for every parallel search_papers instance to complete—no special configuration needed. However, defer=True is useful when branches have different lengths and you want to delay a node’s execution until all other pending tasks are completed Langchain. It’s passed as a parameter to add_node():

Without defer=True, when branches execute in the same superstep, the downstream node naturally waits for all parallel branches to finish Langchain. With defer=True, the node explicitly waits until all pending tasks across different branch lengths are completed—ensuring complete data collection even when branches have asymmetric execution paths.

Key Takeaways

✓ Parallelization delivers dramatic speedups when tasks are independent: By executing multiple nodes simultaneously rather than sequentially, parallel execution can reduce latency. However, this requires careful state management with reducers to merge concurrent updates and avoid data loss.

✓ Subgraphs enable modular multi-agent architectures with isolated or shared state: Use shared state for seamless communication between agents when information needs to flow freely, or isolated state with explicit transformation functions when agents need privacy and clear boundaries between their internal workings.

✓ Map-reduce with the Send API adapts to variable workloads at runtime: Unlike static parallelization with fixed branches, map-reduce dynamically creates parallel tasks based on runtime conditions. This is essential when the number of tasks varies (e.g., processing 2 vs. 10 query interpretations) but adds complexity to debugging and visualization.

References

[1] LangGraph Documentation - Parallelization and Concurrency, LangChain AI, 2024

[2] LangGraph Documentation - Subgraphs, LangChain AI, 2024

[3] LangGraph Documentation - Map-Reduce with Send API, LangChain AI, 2024

[4] LangGraph Documentation - Node Execution and Supersteps, LangChain AI, 2024

This article is also available as a podcast! If you’re on the go or just want to absorb the content in audio format, you can listen to the full episode below 👇 The podcast is also available on Spotify and Apple Podcasts.

As AI agents become more autonomous, chaining tool calls, making decisions, and executing multi-step workflows, the question of how to maintain meaningful human control without sacrificing automation becomes increasingly urgent.

Traditional patterns often force a false choice between fully unsupervised agents, which risk errors and unsafe behavior, and fully supervised workflows, which defeat the purpose of automation.

This tension is especially significant in high-stakes domains such as medical analysis, financial reporting, or content moderation, where an autonomous agent may commit a critical mistake before anyone intervenes, while excessive oversight eliminates efficiency gains.

What is needed is temporal control: the ability to observe reasoning as it happens, pause execution at key moments, inject corrective guidance, and even rewind to earlier states to explore alternative trajectories.

Objective

This article demonstrates how to transform agents from opaque black boxes into steerable, inspectable systems through three complementary mechanisms: streaming for real-time visibility, breakpoints for surgical intervention, and time travel for historical navigation. It relies on LangGraph framework.

These aren’t just debugging features but fundamental primitives for building production-grade agents that balance autonomy with accountability.

After reading this article, you will understand:

1. Streaming: How to expose an agent’s reasoning process in real time, surfacing intermediate thoughts, tool calls, and state transitions as they happen, not after the fact

2. Human-in-the-Loop: How to strategically insert breakpoints that pause execution at critical decision points, allowing humans to inspect context, approve/reject actions, or inject corrective guidance before the agent proceeds

3. Time Travel: How to rewind an agent’s execution to any prior state, fork alternative reasoning paths, and replay modified trajectories, enabling “what-if” exploration and recovery from mistakes without restarting from scratch

Context: This article is the third part of the series on building production-ready AI agents with LangGraph. Part 1 covered the foundational architecture of stateful agents and graph-based workflows, while Part 2 examined persistence mechanisms for checkpoint management and recovery. This new article addresses the most challenging aspect: maintaining human oversight at scale.

Prerequisites:

• Basic knowledge of Python, LLMs, and prompt engineering is recommended.

• Readers should either be familiar with LangChain concepts (chains, models, prompts) or have read Part 1 of this series (or listen the podcast edition) to understand the foundations of LangGraph’s graph, node, and state model.

• Readers should have read Part 2 (or listen the podcast edition) on checkpoint management, state serialization, and recovery mechanisms as these concepts are essential for streaming and time travel.

Tools & libraries: LangGraph, LangChain, LangSmith, OpenAI, arXiv API

You can find the code here on GitHub.

Subscribe now

1. Streaming: Real-Time Visibility into Agent Reasoning

Imagine launching your scientific paper explorer agent with a complex query, then staring at a blank screen for 30 seconds. Is it stuck? Is it searching? Did it crash? Without visibility into what’s happening, you’re forced to trust blindly or worse, kill the process and start over.

This is where streaming transforms your agent from an opaque black box into a transparent, observable system.

1.1 Why Does Streaming Matter for Human-in-the-Loop Systems?

Streaming isn’t just about user experience, though watching tokens appear is certainly more engaging than staring at a loading spinner. For production AI agents, streaming serves three critical functions for human oversight:

1. Trust through transparency: When users see the agent’s reasoning unfold in real-time (the clarified query being formulated, relevance scores being calculated), they can build confidence in the system’s logic or catch problems early.

2. Early intervention signals: If a user watches the agent starting to generate an off-topic query or score irrelevant papers highly, he can interrupt and correct it before wasting compute on the downstream summarization step.

3. Debuggability: When things go wrong, streaming logs capture exactly what the agent was doing at each millisecond. Instead of “the agent failed at minute 2,” you know “it failed while scoring the 8th paper, right after processing a malformed ArXiv result.”

1.2 How to Stream Agent Execution in LangGraph?

LangGraph offers three distinct streaming modes, each revealing different aspects of agent execution:

• Token-Level Streaming: Streams tokens in real time as the LLM generates text.

• Node-Level Streaming: Emits state updates whenever a node finishes running.

• Custom Streaming: Sends progress signals from long-running steps (e.g., API calls, file processing, search loops).

Let’s implement all three in our research assistant to create complete observability.

1. Token-Level Streaming

Key idea: Individual tokens appear as the LLM generates responses, creating a “typewriter effect.” This is useful when you want users to see the summary being written in real-time, or to verify that the LLM is actively generating (not stuck).

Implementation: In the original synchronous summarizer_node, the LLM returns the full summary only once the entire generation is complete. Users see no intermediate output and must wait 10–30 seconds with no feedback while the model works.

The streaming version uses astream() to yield tokens as they’re generated:

The user can progressively see the LLM output with this code:

Output example:

2. Node-Level Streaming

Key idea: It allows users to visualize state updates after each node completes, such as “Clarifier finished” or “Found 5 papers.” This is useful to show high-level progress through the agent pipeline and display intermediate results.

Implementation: Without streaming, only final results are visible.

There are no major code changes needed in the nodes themselves, as LangGraph automatically emits state updates when using stream_mode=”updates”:

Output example:

3. Custom Streaming

Key idea: It displays custom progress messages during slow operations such as API calls (e.g., “Searching ArXiv...”, “Found paper 3/5”). This is particularly useful when a single node performs multiple steps or long-running API requests, allowing for more granular progress updates.

Implementation: The original ArXiv search is a blocking step that provides no intermediate feedback while the query is being executed.

By adding get_stream_writer(), the agent can emit custom progress updates during execution.

This code allows users to see the custom progress updates implemented above.

Combining All Streaming Modes

LangGraph allows multiple streaming modes to be combined simultaneously for complete observability.

Output example:

LangGraph Streaming Modes: Quick Recap

LangGraph provides four streaming modes for different levels of visibility into your agent’s execution:

• “values” streams the complete state after each node. It’s comprehensive but verbose as it shows everything, including unchanged fields.

• “updates” streams only what changed, returned as {node_name: state_updates} It’s the most efficient way to track pipeline progress without noise.

• “messages” streams individual LLM tokens as they’re generated for that typewriter effect. It requires llm.astream() in your nodes and returns (message_chunk, metadata) tuples.

• “custom” lets you emit your own progress signals via get_stream_writer() . It’s perfect for long-running operations like API calls where you want to show “Found paper 3/5” or similar updates.

Combining modes with stream_mode=[”updates”, “messages”, “custom”] is the best way to get multi-level observability. Events then arrive as (mode_name, data) tuples.

2. Human-in-the-Loop: Intercepting and Modifying Agent Execution

Streaming shows users what the agent is doing, but what if you want to let them actually intervene? What if the clarifier misinterprets the query, the researcher ranks irrelevant papers too highly, or you want users to manually approve expensive API calls before they run?

This is where LangGraph’s interrupt system turns agents from autonomous black boxes into collaborative tools that can pause mid-execution, wait for human input, and incorporate corrections before continuing.

Interrupts enable three essential capabilities:

1. Approval workflows: Pause before expensive or sensitive operations (like sending emails, making purchases, or calling paid APIs) and require explicit human confirmation to proceed.

2. Edits and corrections: Stop execution when the agent generates intermediate results (refined queries, paper rankings, draft messages), let the user modify the state, then resume with the corrected inputs.

3. Debugging and inspection: Freeze execution at specific nodes to inspect state, inject test data, or verify that complex logic behaves as intended before continuing.

2.1 How Does LangGraph Interrupt System Work?

LangGraph provides two complementary ways to introduce controlled pauses during execution:

1. Graph-level interrupts: These let you designate specific nodes as pause points when building your workflow. Whenever execution reaches one of these points, the graph stops automatically. You can then review the intermediate state, make adjustments, and resume the workflow from the exact same spot.

2. Node-level interrupts: These allow a node itself to request input mid-execution. The node temporarily hands control back to the user, waits for the missing information or correction, and then continues once the needed value is provided.

In both cases, LangGraph relies on its checkpointer to capture the full execution state. This ensures the pause and the subsequent resume feel seamless and reliable.

In our scientific paper explorer, interrupts eliminate a classic frustration: watching the agent spend 30 seconds summarizing papers you already know are irrelevant just from reading their titles. Let’s implement human-in-the-loop controls at three strategic points in our pipeline.

1. Graph-Level Interrupts

Key idea: This approach introduces explicit pause points directly in the graph definition. Execution automatically stops at these points, and interaction happens externally through get_state() and update_state(). It enables reviewing intermediate outputs, adjusting state based on observations, or inserting approval steps, all without modifying node logic.

Implementation: The pattern looks like this:

In the scientific paper explorer, cost awareness matters. Before the researcher begins querying the ArXiv API and generating expensive summaries, it can be useful to inspect the clarifier’s refined query. With breakpoints, the workflow pauses at that moment, allowing review or adjustment of the query before the agent continues.

This example highlights how breakpoints work in practice:

• The graph is compiled with interrupt_before=[”researcher”] and interrupt_after=[”clarifier”], causing execution to pause right after the clarifier and before any expensive API calls.

• A thread_id in the config lets LangGraph store and retrieve the execution state via the checkpointer.

• The first graph.stream(...) runs until an interrupt is reached, then stops.

• graph.get_state(config) retrieves the paused state (including the refined query).

• You can modify this state with graph.update_state(...) or leave it unchanged.

• Calling graph.stream(None, config, ...) resumes execution from the checkpoint using the updated state.

This creates a tight approval loop that gives you control before costly steps execute.

2. Node-Level Interrupts

Key idea: Graph-level interrupts are great when you can define fixed pause points ahead of time, but what if you want a more user-friendly way to request and handle user input directly within a node’s logic? What if a node needs to ask for input dynamically based on what it discovers at runtime? That’s where node-level interrupts shine.

Implementation: Let’s create a dedicated approval node that pauses execution and requests a human decision:

Once the approver node is in place, you still need a conditional edge so that if approval is denied, the workflow routes directly to END instead of continuing through the remaining nodes.

This modifies the structure of the agent graph: after adding the approver node, the workflow now branches based on the approval result, as shown below.

The example above shows how a node can pause itself, request user input, and continue with that input. Technically, this works as follows:

• The interrupt() call inside the approval node pauses execution at that exact line and stores the interrupt request in the checkpointer under the given thread_id.

• The paused interrupt payload appears in state.tasks, which is why the code retrieves it to display the message to the user.

• When the user responds, the value is passed back into the graph via Command(resume=...), which feeds the response directly into the waiting node.

• Execution then resumes right after the interrupt, and the node updates the state based on the user’s decision.

• A conditional edge reads the approved flag to route the workflow either forward to the researcher or directly to END.

This creates a fully interactive pause-and-resume node driven by user input.

3. Time Travel: Rewinding and Replaying Agent Execution

Interrupts let you pause and modify your agent’s future, but what if you need to change its past? What if you realize three nodes later that the clarifier’s query refinement was wrong, or you want to A/B test different researcher parameters without starting from scratch?

This is where LangGraph’s time-traveling features excel. Building on the same checkpointing system that enables memory and interrupts, time travel lets you rewind execution to any previous state, modify it, and create alternate execution branches, essentially giving you “undo” and “what-if” superpowers for your agent workflows.

Time travel serves three powerful use cases:

1. Rewinding mistakes: When you discover an error several steps deep (like a misunderstood query leading to irrelevant papers), jump back to that decision point, fix it, and re-run from there without losing subsequent work.

2. A/B testing and experimentation: Fork execution at a checkpoint to try different parameters (e.g., “what if I used 10 papers instead of 5?”) and compare outcomes without re-running the entire pipeline.

3. Debugging and root cause analysis: Step backward through execution history to pinpoint exactly where things went wrong, examining state at each checkpoint like a debugger’s stack trace.

3.1 How Does Time Travel Work in LangGraph?

LangGraph’s time travel relies on the checkpointer, which doesn’t just store the latest state, but maintains a complete history of every state transition. Each checkpoint is identified by a unique ID, forming a branching tree of execution states. You can navigate this tree like Git commits: inspect any historical state, rewind to it, and optionally fork a new branch from that point.

The key methods are:

• graph.get_state_history(config): Retrieve all checkpoints for a thread, from newest to oldest

• graph.get_state(config, checkpoint_id=id): Jump to a specific checkpoint

• graph.update_state(config, values, as_node=”node_name”): Rewind to before a specific node and replay with modified state

Rewinding to Fix Mistakes

Imagine you’ve run clarification, paper search, and scoring, only to realize during summarization that the query refinement missed the mark. Instead of starting over, rewind to the clarifier, fix it, and fast-forward:

Time travel is possible because of three core mechanisms:

1. State is saved after every node: Each node execution produces a new state snapshot stored in the checkpointer with a unique identifier.

2. Checkpoints form a branching tree: Editing a past checkpoint and resuming execution creates a new branch, leaving the original path untouched.

3. The as_node parameter is crucial: It instructs LangGraph to treat the checkpoint as if a specific node just executed, effectively rewinding the graph to that point.

Initial run:
  start → clarifier → researcher → summarizer ✓

After rewind with as_node=”clarifier”:
  start → clarifier ─┬→ researcher (old) → summarizer (old)
                     └→ researcher (new) → summarizer (new) ✓

The original branch remains intact in the history; modifying a checkpoint simply creates an alternate timeline.

Creating Parallel Branches for A/B Testing

Time travel goes far beyond simple corrections, it enables creating multiple branches from the same checkpoint, allowing different approaches to be compared without repeating costly steps.

For example, suppose you want to test whether fetching 5 papers or 10 papers yields better summaries. Instead of running the clarifier twice (an expensive LLM call), run it once, then fork two branches from that checkpoint with different max_papers values. Both branches reuse the clarifier’s output and diverge only at the researcher node, enabling an efficient A/B test while paying the clarification cost only once.

By assigning different thread_id values to each branch, you maintain fully independent execution paths that can be evaluated side by side.

Initial run (thread: base):
  start → clarifier → researcher(5) → summarizer ✓

After creating branches:
  start → clarifier ─┬→ [thread: base]     researcher(5)  → summarizer
                     ├→ [thread: branch-a] researcher(5)  → summarizer  
                     └→ [thread: branch-b] researcher(10) → summarizer
All three branches share the clarifier checkpoint but diverge afterward.

With time travel, your agent becomes fully explorable and correctable. You can undo mistakes without re-running expensive operations, compare alternative execution paths side-by-side, and debug complex workflows by stepping through history like a version control system. It’s the ultimate safety net for production AI agents—and the foundation for letting humans collaborate with agents iteratively rather than hoping they get it right on the first try.

Key Takeaways

✓ Streaming provides three complementary layers of agent observability: Token-level streaming reveals LLM generation as it happens, node-level streaming exposes state transitions between workflow steps, and custom streaming surfaces progress during long-running operations. Together, they transform opaque agents into transparent, debuggable systems.

✓ Human-in-the-loop patterns balance automation with control through strategic pause points: Graph-level interrupts halt execution at predefined nodes for review and approval, while node-level interrupts allow dynamic requests for human input mid-execution. Both maintain full context through checkpointing, enabling seamless resume without data loss.

✓ Time travel enables non-destructive experimentation and error recovery: Checkpointers create a complete execution history where any prior state can be inspected, modified, and replayed. Mistakes become fixable through rewinding rather than restarting, and parallel branches allow A/B testing without duplicating expensive operations.

References

[1] LangChain Documentation. (2025). LangGraph: Streaming modes for real-time agent observability.

[2] LangChain Documentation. (2025). LangGraph: Human-in-the-loop with interrupts and breakpoints.

[3] LangChain Documentation. (2025). LangGraph: Persistence and checkpointers for stateful agents.

[4] LangChain Documentation. (2025). LangGraph: Time travel and state history navigation.

[5] OpenAI Documentation. (2025). Streaming completions with Server-Sent Events.

[6] arXiv Documentation. (2025). arXiv API User’s Manual: Query interface and metadata retrieval.

[7] LangSmith Documentation. (2025). LangSmith: Tracing and debugging LangGraph workflows.

This article is also available as a podcast! If you’re on the go or just want to absorb the content in audio format, you can listen to the full episode below 👇 The podcast is also available on Spotify and Apple Podcasts.

Building a demo agent is easy. Chain a few LLM calls, add some tools, and watch it work. But as your agent grows more capable, it becomes more fragile. More steps mean more failure points. More state means more memory pressure. More iterations mean higher costs from redundant work.

This is the paradox of agentic complexity: the more powerful your agent, the more likely it can fail catastrophically.

The gap between prototype and production is persistence. Without it, agents remain fragile toys, impressive in demos, unreliable in reality. To bridge this gap, three core mechanisms bring production-grade durability to AI systems:

• Reducers and caching to optimize state updates and skip redundant computation

• Persistence and memory to preserve context across sessions

• Checkpointers to recover from interruptions without losing progress

Reducers keep state manageable as workflows lengthen. Memory prevents exponential cost growth. Persistence scales with complexity. Together, they turn fragile reasoning graphs into robust, resumable intelligence pipelines.

The result: agents where capability and reliability grow together, not inversely.

Objective

This article focuses on implementing persistence mechanisms in LangGraph, from basic memory to advanced checkpointing, to build agents that remember context, optimize computation, and survive failures.

After reading this article, you will understand:

1. How to optimize state with reducers and caching — managing growing state efficiently through composable update functions and skipping expensive recomputation through intelligent caching strategies.

2. How to implement persistence and memory — maintaining state across sessions, preserving conversation history, and ensuring agents remember what they’ve already done to avoid redundant work

3. How to build fault-tolerant systems with checkpointers — saving state at every step, resuming execution from any point, and recovering gracefully from failures without losing progress

Context: This article is Part 2 of a series on building production-grade agents with LangGraph. Future articles will explore advanced multi-agent orchestration patterns, deployment strategies, and scaling considerations for complex agentic systems.

Prerequisites: Basic knowledge of Python, LLMs, and prompt engineering is recommended. Readers should either be familiar with LangChain concepts (chains, models, prompts) or have read Part 1 of this series (or listen the podcast edition) to understand the foundations of LangGraph’s graph, node, and state model.

Tools & libraries: LangGraph, LangChain, LangSmith, OpenAI, arXiv API

You can find the code here on GitHub.

Subscribe now

1. Reducers & Cache: Doing More with Less

Every interaction with your agent adds weight to its state. Messages accumulate. Retrieved documents pile up. Tool outputs compound. Without management, a research agent handling several queries consumes megabytes of context, triggers token limit errors, and starts lagging.

The solution isn’t storing less, it’s storing smarter.

Two mechanisms help overcome this issue:

• Reducers trim state growth by keeping only what matters.

• Caches that skip redundant computation.

Used together, they turn memory from a liability into an asset that scales efficiently with your agent’s complexity.

1.1 How Do States Get Updated?

Let’s revisit the scientific paper explorer from Part 1. This agent consists of three specialized nodes.

The graph includes intelligent retry logic: if fewer than 3 papers are found, it automatically loops back to search again (up to a configurable max iteration limit).

Here’s the state definition:

Most fields are straightforward (str, int, List[dict]), but messages includes something special: Annotated[List[BaseMessage], add]. That little annotation changes everything; it defines how state updates behave.

The Default: Complete Replacement

• By default, LangGraph replaces values when updating state.

• This is perfect for scalar fields like query or summary: as one would only need the latest version.

• But for collections like messages or papers, replacement is disastrous. Each node would overwrite prior data, erasing the conversation history, previous research results, or user inputs. The agent would lose its memory every time it updates.

The Alternative: Reducers

• Reducers are functions that control accumulation, deduplication, or any other merge logic.

• In the scientific agent example: messages: Annotated[List[BaseMessage], add]tells LangGraph: “When updating messages, don’t replace. Use the add function to combine old and new values.”

• The add operator (from Python’s operator module) concatenates lists, allowing the conversation to build naturally.

1.2 How Do Reducers Control Memory Growth?

Preserving history is good, but endless accumulation is not. As an agent reviews more papers, each interaction adds new messages until memory overflows. Context length explodes, inference slows, and token costs rise. This is where reducers shine: they can control how memory grows, not just preserve it.

Let’s explore three frequently used reduction strategies:

1. Truncation Reducers

   This is typically used when only recent context matters and historical data becomes irrelevant. For instance, a real-time monitoring agent tracking system errors only needs the last 20 events. Last week’s logs don’t inform today’s diagnosis.
   Trade-offs:

   ✅ Simple and predictable memory footprint.

   ❌ The agent completely forgets history, losing the ability to reference or learn from past patterns beyond your window size.

   

2. Selective Retention

   This is ideal when you’re collecting data with varying quality and only want to keep the best items. Think of a research agent that finds 50 papers but only needs the 10 most relevant ones. Instead of keeping everything, only the most valuable is preserved.
   Trade-offs:
   ✅ Maintains bounded memory while keeping the highest-quality information.
   ❌ Requires domain-specific logic to define “quality, with the risk of permanently losing items that seem low-value now but could be relevant later.

   

3. Summarization Reducers

   This is perfect when context is valuable but verbosity isn’t. For instance, a research agent should remember what papers it’s seen and why, but doesn’t need the full back-and-forth discussion. Summarization compresses old context while preserving recent details.

   Trade-offs:

   ✅ Balances memory efficiency with context preservation.

   ❌ Loses fine-grained details in compressed sections, and summarization itself costs tokens (though far fewer than keeping everything)

   

   For the paper explorer, selective retention works best for papers while summarization makes the most sense for messages. This hybrid approach ensures both bounded memory and preserved context where it matters. You can find the complete implementation in the GitHub repository.

1.3 How Does Cache Avoid Redundant LLM Calls?

Reducers solve memory growth, but there’s a complementary problem: computational waste. Agents often repeat the same operations, re-ranking identical papers, re-answering similar questions, and re-processing unchanged inputs. Every repetition burns tokens and time.

Caching captures the results of expensive operations and reuses them when inputs match.

For the paper explorer, this means massive savings:

• Clarifying similar queries: With a semantic cache, the agent would be able to return the same refined query with limited costs for queries similar to “Tell me about transformers” like “Explain transformer models”.

• Ranking papers: The cache would return stored relevance scores without re-evaluation for the same paper IDs and queries.

• Summarizing papers: The cache would return the existing summary for already parsed papers.

Semantic vs. Exact Caching

• Exact caching matches inputs byte-for-byte. It’s fast and reliable but brittle. Examples of mismatches:

  • “Summarize this paper”

  • “Summarize this paper.” # extra period

  • “ Summarize this paper” # extra space

• Semantic caching matches inputs by meaning (for example, using embeddings and cosine similarity). Slightly slower to determine equivalence but far more robust for natural-language variants. Examples of matches:

  • “What are transformers in NLP?”

  • “Explain transformer architectures for natural language”

  • “Tell me about attention mechanisms in language models”

Inside the Cache

LangChain provides built-in caching at the LLM level through cache backends. When you make an LLM call, LangChain:

1. Generates a cache key from the prompt, model parameters, and model name

2. Checks the cache for an existing result

3. Returns cached response if found, or calls the LLM and stores the result

This process runs seamlessly under the hood, requiring no code modifications.

LangChain offers several cache backends, each with different tradeoffs in persistence, scalability, and setup complexity.

1. In-memory cache (development)

2. SQLite cache (persistent, local)

3. Redis cache (production-ready)

1. In-Memory Cache (Development)

It’s a cache that lives in RAM. It’s fast and simple to add, but everything vanishes when the app stops.

2. SQLite Cache (Persistent, Single-Instance)

It’s a local database-backed cache that survives restarts. It’s perfect for single-instance tools.

3. Redis Cache (Production, Distributed)

It’s a distributed cache backed by Redis. It’s fast, concurrent, and shareable across servers.

For the paper explorer example, SQLite is a good starting point for development. When deploying multiple instances or sharing cache across users, Redis would be more appropriate.

Semantic Caching in LangChain

For semantic similarity matching, LangChain offers specialized cache implementations:

When a query comes in, LangChain:

1. Embeds the prompt using your chosen embedding model

2. Searches cached embeddings via vector similarity

3. Returns cached response if similarity exceeds the threshold

4. Stores new embedding and response if cache miss

Note: LangChain doesn’t provide built-in in-memory or SQLite semantic caches, but it can easily be implemented.

2. Persistent Dialogues: Threads, State, and Sessions

Reducers and caches optimize individual state updates and computations. But production agents need something more: the ability to maintain coherent conversations across multiple interactions.

When a user asks “What papers did you find?” then follows up with “Summarize the third one,” the agent needs to know which papers were retrieved and which one is “the third.” When a conversation crashes mid-analysis, the user shouldn’t have to start over. When multiple users interact with the same agent simultaneously, their conversations must remain separate and private.

This is conversation continuity: the infrastructure that transforms stateless LLM calls into stateful dialogue systems.

2.1 How Do Thread IDs Isolate Conversations?

Every conversation needs an identity. In LangGraph, this identity is the thread ID, a unique identifier that isolates one conversation from another.

Think of threads like separate browser tabs. Each tab maintains its own history, state, and context. Actions in one tab don’t affect others. The same principle applies to agent conversations:

• User A (thread: user_a_session_123) asks about quantum computing papers

• User B (thread: user_b_session_456) asks about climate change research

• Both conversations run independently, never mixing context or state

Implementation is simple: Pass a thread ID in your configuration when invoking the graph.

What the thread ID enables:

• Isolation: Each conversation maintains separate state, messages, and context.

• Persistence: State survives between invocations, users can leave and return hours later.

• Continuity: Follow-up questions reference earlier messages without explicit context passing.

• Concurrency: Multiple conversations run simultaneously without interference.

Real-world thread ID patterns:

The thread ID doesn’t just label conversations. It’s the key that unlocks every checkpoint, message, and piece of state associated with that dialogue.

2.2 How To Define An Efficient State Schema Design?

Not all states should be visible everywhere. The paper explorer might track internal ranking scores, API rate limits, or retry counts, information crucial for workflow logic but irrelevant to the final output. Exposing everything creates noise. Hiding everything blocks debugging.

LangGraph provides three schema patterns to control state visibility:

1. Public state

2. Private state

1. Public State: The Default

This consists of designing a single schema shared across all nodes and exposed as both input and output. Every node sees and can modify every field. The entire state is returned as output. Simple, but sometimes too simple.

2. Private State: Node-to-Node Secrets

Sometimes nodes need to exchange information that shouldn’t pollute the main schema. Consider the paper explorer workflow:

• Clarifier Node: This node generates an optimized ArXiv search query (refined_query) which is just a technical transformation. For instance, users don’t need to see “transformer attention mechanism” when they ask “how do transformers work?”

• Researcher Node: This node uses the refined query to search, then immediately discards it as the papers are what really matter.

• Summarizer Node: This node only needs the final ranked papers, not the search query that found them

Private state lets nodes pass ephemeral working data without exposing it in the final output.

Let’s build this concept step by step using a paper search workflow.

• Step 1: Define Your Public Interface

  Start by defining what users provide and what they receive:

  

  Key insight: Notice what’s not here: no internal counters, no refined queries, no technical scaffolding. Just the essentials.

  

  Think of InternalState as a “kitchen workspace”, it contains both the ingredients (input), the final dish (output), and all the prep bowls and measuring cups (internal fields).

• Step 2: Build Your Nodes

  Now let’s see how nodes use both public and private state:

  

• Step 3: Connect Everything

  

The user interface is clean and focused:

Private state is perfect for:

• Intermediate transformations that users don’t care about

• Internal counters and flags

• Raw technical data such as API response headers, performance metrics or internal timestamps

They are not useful for:

• Debugging information: It’s better to keep these in public state during development (you can filter them out in production with OutputState)

• Data needed by distant nodes: If a node 5 steps away needs a variable, it’s better to make it public. Private state should be short-lived (1-2 nodes)

• User preferences: Configuration like max_papers or llm_model should be in InternalState for all nodes to access.

3. Checkpointers: Never Losing Progress

Reducers manage memory growth. Caches skip redundant work. Thread IDs isolate conversations. But there’s one failure mode that none of these solve: what happens when something crashes mid-execution?

For instance, your research agent finds 8 relevant papers, summarizes 5 of them, then hits a rate limit. Without checkpointing, all progress will vanish. You must start over, re-searching the same papers, re-generating the same summaries. Every crash means wasted tokens, wasted time, and frustrated users.

Checkpointers eliminate this fragility. They automatically save your agent’s state after every step, creating a progress trail. When failures occur, the agent simply picks up where it left off, no data lost, no work repeated.

3.1 How Does Checkpointing Work?

Think of checkpointers as version control for agent execution. After each node completes, LangGraph captures a snapshot of the entire state and stores it persistently. Each snapshot includes the current values of all state fields, the conversation history, and metadata about which node just executed.

When you invoke the graph with a thread ID, LangGraph:

1. Checks for the existing state associated with that thread

2. Loads the latest checkpoint if one exists

3. Resumes execution from the next node, not from scratch

4. Saves a new checkpoint after every node completes

5. Continues this pattern until the graph reaches an end state

The key insight: With checkpointing, your agent becomes resumable by default.

3.2 How To Implement Checkpointers?

LangGraph provides three checkpointer backends, each trading simplicity for durability:

1. MemorySaver

2. SQLite

3. Postgres

1. MemorySaver: Development and Testing

The simplest option stores checkpoints in RAM. It’s perfect for local development and testing, but everything disappears when the process stops.

2. SQLite: Single-Instance Persistence

SQLite provides durable storage with zero infrastructure. Checkpoints survive restarts, making it ideal for single-instance deployments, personal tools, or applications where multiple servers don’t need to share state.

3. Postgres: Production-Scale Persistence

For production systems with multiple instances, concurrent users, or distributed architectures, Postgres provides the durability and concurrency control needed. Multiple servers can safely read and write checkpoints without conflicts.

3.3 How To Resume from Failure?

The real power of checkpointing emerges when things go wrong. With checkpoints enabled, your agent can recover from failures automatically by loading the last successful state and continuing execution.

Example: Research agent hits rate limit

What happened internally:

• First run: Clarifier ran, researcher found 8 papers, summarizer processed 5 papers (checkpoints saved after each node)

• Crash: Rate limit hit while processing paper 6

• Second run: Loaded checkpoint after paper 5 was summarized, resumed with paper 6, no repeated API calls or LLM invocations

This pattern extends beyond crashes. Checkpointing enables human-in-the-loop workflows where agents pause for approval, long-running tasks that span hours or days, and multi-session conversations that persist indefinitely. More on this in the coming articles.

How To Debug with Checkpoints?

Checkpointers don’t just enable recovery, they create a complete execution history. You can inspect any previous state, understand how values changed, and debug issues by replaying the exact sequence of steps.

This visibility is crucial for:

• Understanding why the agent made specific decisions

• Identifying which step introduced incorrect data

• Comparing different execution paths for the same query

• Building audit trails for compliance or debugging

Key Takeaways

✓ Reducers prevent memory overflow by controlling how state accumulates: Instead of endless growth, use truncation (keep last N items), selective retention (keep only highest-quality items), or summarization (compress old context) to maintain bounded memory while preserving essential context.

✓ Semantic caching eliminates redundant LLM calls by matching intent, not exact strings: Unlike exact caching that fails on trivial variations, semantic caching uses embeddings to recognize equivalent queries across different phrasings, dramatically reducing costs for repetitive operations.

✓ Thread IDs and checkpointers transform fragile demos into production-ready agents: Thread IDs isolate concurrent conversations while checkpointers save state after every node, enabling automatic recovery from failures without losing progress or repeating expensive work.

✓ Private state keeps interfaces clean by hiding implementation details: Separate InputState (user provides), OutputState (user receives), and InternalState (nodes share internally) to pass technical data between nodes without polluting final output.

References

[1] LangChain Documentation. (2025). LangGraph: Build resilient language agents as graphs.

[2] LangChain Documentation. (2025). LangChain: Building applications with LLMs through composability.

[3] LangSmith Documentation. (2025). LangSmith: Platform for debugging, testing, evaluating, and monitoring LLM applications.

[4] LangGraph Documentation. (2025). Persistence and Memory: Checkpointer implementations.

[5] LangChain Documentation. (2025). Caching: LLM response caching strategies.

[6] OpenAI Documentation. (2025). OpenAI Platform.

[7] arXiv Documentation. (2025). arXiv API User’s Manual.

This article is also available as a podcast! If you’re on the go or just want to absorb the content in audio format, you can listen to the full episode below 👇 The podcast is also available on Spotify and Apple Podcasts.

Consider an agent that answers complex scientific questions by retrieving papers from arXiv, cross-referencing findings, and synthesizing insights. It works perfectly, until an API fails, a token limit hits, or reasoning loops infinitely.

Linear chains can’t recover. No loops to retry. No branching to adapt. No state to resume from.

LangGraph changes this. It models workflows as stateful graphs where each node is a reasoning step. State persists. Control flow adapts. Failures become recoverable.

The result? Agents that don’t just execute, they reason iteratively, remember context across steps, and recover gracefully. Exactly what’s needed when the path forward isn’t predetermined.

Objective

This article focuses on building stateful, graph-based AI agents with LangGraph by progressively constructing a scientific research assistant that retrieves, analyzes, and synthesizes academic papers.

After reading this article, you will understand:

• How to build graph architectures from the ground up—defining state schemas, creating node functions, and connecting them with edges to enable workflows that sequential chains cannot handle

• How to add intelligence through routing and agents—implementing conditional edges for dynamic decisions, binding tools for autonomous actions, and transforming simple chains into agentic nodes

• How to implement persistence and recovery patterns—using memory to prevent redundant work, checkpointing to resume from failures, and streaming to monitor agent execution in real-time

Context: This article is Part 1 of a series on building production-grade agents with LangGraph. Future articles will explore advanced multi-agent orchestration patterns, deployment strategies, and scaling considerations for complex agentic systems.

Prerequisites: Python programming experience and basic understanding of LLMs and prompt engineering. Familiarity with LangChain concepts (chains, models, prompts) is NOT required.

Tools & libraries: LangGraph, LangChain, LangSmith, OpenAI, arXiv API

You can find the code here on GitHub.

Preview: Below is a demo of the graph-based research agent described in this article, showcasing iterative search, reasoning, and synthesis with proper citations.

Subscribe now

1. The Research Problem: Why Linear Agents Can’t Think Like Scientists

Before exploring graph-based agentic systems, let’s examine the problem that will serve as our guiding example throughout this article.

Imagine you want to build an AI research assistant capable of answering scientific questions by consulting actual research papers, retrieving them, analyzing their contents, and synthesizing evidence-based answers. A typical question would be: “What are the recent advancements in using CRISPR for cancer immunotherapy?”

At first, this seems simple enough for a single linear chain: query an API, summarize the papers, and generate an answer. Yet, scientific reasoning rarely follows a straight path. What happens when key references appear in papers not yet retrieved? When early results are too superficial, demanding refined searches? Or when new findings open unexpected research directions? A linear chain cannot handle these situations.

Graph-based reasoning solves this by reflecting how researchers actually think and work:

• Non-linear exploration: It can search again after reading papers that reveal new directions.

• Conditional decisions: It can evaluate paper quality and determine whether to search more or synthesize findings.

• Iterative refinement: It can loop back to search for specific details that emerged during analysis.

• State-aware branching: It can adapt the strategy based on the question type, with some requiring breadth and others demanding depth.

The Practical Implementation

The scientific paper explorer presented in this article serves as a concise yet complete example of a graph-based research agent. It consists of three specialized nodes, each designed to perform a distinct cognitive function.

To ensure robustness, the graph incorporates intelligent retry logic. When the researcher node retrieves fewer than three papers, the system automatically loops back to refine the query and search again, up to a configurable maximum iteration limit.

This simple structure already demonstrates the advantage of graph-based reasoning: the ability to adapt, revisit steps, and iteratively improve results without relying on a rigid linear chain.

2. Graph Architecture: Designing the Backbone of Adaptive AI Agents

2.1 What Defines the Agent’s State?

Every agent needs memory: a way to carry context from one reasoning step to another. In LangGraph, this is represented as a state object, often implemented through a TypedDict that defines what kind of data flows through the graph.

TypedDict is a Python feature (from the typing module) that lets you define dictionary types with specific keys and value types, providing type hints for dictionaries with a fixed structure.

For our AI research assistant, the state is defined as follows:

This state evolves through the reasoning process:

• query: The user’s original scientific question

• refined_query: The optimized version used for the search

• papers: The list of retrieved results (metadata, abstracts, etc.)

• summary: The structured, synthesized output

• iteration: A safeguard against infinite loops

• messages: The full conversational trace (LangChain messages)

Together, these elements capture both content and context, enabling the agent to reason iteratively while preserving history for monitoring and debugging.

Notice the Annotated[List[BaseMessage], add] syntax on messages. This is a reducer annotation that tells LangGraph how to handle updates. Without a reducer, returning {”messages”: [new_message]} would replace the entire conversation history. But with the add reducer, it automatically appends new messages to the existing list, building up the conversation naturally without each node needing to manually read and concatenate. More on this in the comming articles.

2.2 How to Structure Tasks into Nodes?

Once the state is defined, the next step is to design the nodes—the reasoning units that transform that state.

In LangGraph, each node is a Python function performing a distinct cognitive task. Think of them as the agent’s mental operations: clarifying a question, searching for information, analyzing findings, and deciding what to do next.

Each node is a function that:

1. Takes the current AgentState

2. Performs a reasoning step (e.g., refine query, fetch data, synthesize results)

3. Returns an updated version of the state

This modular approach makes the workflow transparent, debuggable, and flexible.

Our research agent includes three main nodes:

1. clarifier_node — Refines the original question into a precise query

2. researcher_node — Uses an ArXiv tool to fetch relevant papers

3. summarizer_node — Summarizes the papers and decides whether to continue

Here is an example of a node: the clarifier_node node

This node extracts the user’s question from state, configures an LLM with specified parameters, and uses a specialized system prompt to transform natural language into precise ArXiv search terms. The refined query is then stored back into state along with the conversation history for downstream nodes.

Each node has a clear, single responsibility:

• Clarifier: Transforms natural language into precise search terms

• Researcher: Handles external API interactions and data retrieval

• Summarizer: Analyzes content and makes strategic decisions about next steps

The beauty of this separation is modularity. Want to search PubMed instead of ArXiv? Replace the researcher_node. Need an additional quality-check step? Insert a new node between researcher and summarizer. The shared state ensures compatibility without rewriting the logic.

Together, these nodes form a self-consistent reasoning pipeline where every step is visible and testable.

2.3 How to Connect Nodes into a Graph?

With nodes defined, the final step involves linking them into a graph that defines the agent’s reasoning flow.

LangGraph uses a StateGraph to specify:

• Nodes: The reasoning units

• Edges: Control flow between operations

• Conditional edges: Logic-based transitions (looping back if too few papers are found)

This creates a structured yet flexible reasoning loop, ideal for research workflows where iterative exploration is essential.

For our AI research assistant, the graph is defined as follows:

Graph construction proceeds in three steps: registering nodes, establishing connections, and implementing conditional logic.

1. Registering Nodes: First, each processing unit is registered as a named node in the workflow, representing a distinct phase in the agent’s pipeline: workflow.add_node(”clarifier”, clarifier_node) workflow.add_node(”researcher”, researcher_node) workflow.add_node(”summarizer”, summarizer_node)

2. Establishing Flow: Next, the execution path is defined. The clarifier serves as the entry point, with direct edges chaining the nodes together: workflow.set_entry_point(”clarifier”)

   workflow.add_edge(”clarifier”, “researcher”)

   workflow.add_edge(”researcher”, “summarizer”)
This creates a linear progression where each stage feeds into the next.

3. Implementing Conditional Logic: The critical intelligence comes from conditional edges. After the summarizer completes its analysis, the should_continue function determines the next action:
   workflow.add_conditional_edges(

   “summarizer”,

   should_continue,

   {“researcher”: “researcher”,

   “end”: END}

   )

   This establishes a feedback loop: when additional papers are needed, execution returns to the researcher node; otherwise, the workflow terminates.

4. Calling workflow.compile() produces an executable graph ready to process research queries with adaptive iteration based on result quality.

Note: The checkpointer enables persistence, allowing the graph to save state snapshots at each node execution. This supports advanced features like pausing execution, resuming from specific points, or implementing human-in-the-loop workflows. Future articles will explore these capabilities in depth.

3. Tracing, Debugging, and Monitoring: The Path to Reliable AI Agents

3.1 Why Agent Development Requires Dedicated Tooling?

Building agents with LLMs is fundamentally different from traditional software development. The challenges that make specialized tooling essential include:

• Non-deterministic behavior: LLMs are non-deterministic by nature, producing varied outputs even for similar inputs, making debugging and quality assurance more complex than traditional software. A prompt that works perfectly in testing might fail mysteriously in production. Without proper observability, developers are left guessing why an agent chose a particular path or generated unexpected output.

• Multi-step reasoning opacity: The opaque nature of LLMs can make debugging, monitoring, and performance optimization a formidable task. When an agent fails, the question isn’t just “what went wrong?” but “at which step in the reasoning process did it derail?”. Traditional debuggers with breakpoints can’t capture the cognitive flow of an agent that searches papers, analyzes findings, and decides whether to search again.

• State complexity across nodes: Agents traverse multiple nodes with intermediate states that need inspection. In the research agent example, the state evolves from an empty papers list, to populated papers, to extracted findings, to a final summary. Each transformation represents a potential failure point. Without visibility into state transitions, debugging becomes impossibly difficult.

• Iterative loops and conditional branching: Graph-based agents can loop back to previous nodes or branch based on conditions. The agent can interrupt at any time, run in debug mode where it pauses after each step, and developers can walk through step-by-step. Traditional linear debugging tools weren’t designed for workflows that might execute “search → analyze → search again → analyze → synthesize.”

• Production monitoring challenges: In production, there’s significantly more traffic, making it impractical to examine datapoints individually. Developers need aggregate metrics, anomaly detection, and the ability to drill down from high-level dashboards into specific problematic traces.

These challenges demand tools purpose-built for agentic AI—tools that understand graph structures, state management, LLM calls, and the unique debugging needs of multi-step reasoning systems.

3.2 How Does LangGraph Studio Enable Rapid Prototyping and Debugging?

LangGraph Studio offers a new way to develop LLM applications, providing a specialized agent IDE for visualizing, interacting with, and debugging complex agentic applications.

Visual Graph Representation

Design workflows in a flowchart-style editor by connecting nodes and edges. The graph updates in real time as code changes, so structural issues become obvious: missing edges, incorrect routing, or unintended cycles that would be easy to miss in code.

Interactive State Inspection

At any point, developers can interact with the state of the agent; if they don’t like what the agent responded with at a specific step, they can directly modify the response and then continue with that new response.

This is transformative for debugging: imagine the summarizer decides it needs more papers even though 5 papers were found.

• Traditional debugging: Add print statements, restart the entire flow, hope to catch the issue.

• Studio debugging:

  • Pause execution at the summarizer node

  • Inspect the state: {”papers”: [5 papers], “iteration”: 1, ...}

  • Notice the issue: papers have low relevance scores

  • Edit the state to inject better papers

  • Resume execution to test the hypothesis

This “time travel” debugging capability is invaluable for understanding agent behavior without repeatedly running expensive LLM calls.

Real-Time Code Hot Reloading

LangGraph Studio detects changes to the underlying code files, allowing developers to update prompts in their code editor and rerun nodes if an agent responds poorly, making it much easier to iterate on long-running agents. The workflow becomes:

1. Run the agent in Studio

2. Notice that the clarifier node produces a poor query refinement

3. Edit the system prompt in nodes.py

4. Studio detects the change and hot-reloads

5. Rerun just the clarifier node with the same input

6. Instantly see if the new prompt improves results

No container rebuilds, no full restarts—just immediate iteration on prompts, logic, and LLM parameters.

3.3 How Does LangSmith Provide Production Observability and Monitoring?

LangSmith Observability gives complete visibility into agent behavior with tracing, real-time monitoring, alerting, and high-level insights into usage.

While Studio excels at local development, LangSmith is purpose-built for production environments where scale, reliability, and performance matter.

To get started with tracing, you only need to set one environment variable if already using LangChain or LangGraph.

For the research agent:

LANGCHAIN_TRACING_V2=true
LANGCHAIN_API_KEY=”XXX”
LANGCHAIN_PROJECT=scientific-graph-agent

End-to-End Trace Inspection

Every invocation of the graph automatically generates a detailed trace capturing the entire reasoning flow, including the initial user query, each node execution (clarifier, researcher, summarizer), all LLM prompts and responses, tool calls such as ArXiv API searches, intermediate state updates, conditional edge decisions, and the final output with total execution time.

A trace represents the complete end-to-end execution of a request, capturing every significant operation as a run within that trace.

For the research agent, a single trace might look like:

Developers can quickly debug and understand non-deterministic LLM app behavior with tracing, seeing what the agent is doing step by step, then fixing issues to improve latency and response quality.

Production Monitoring Dashboards

In production, with significantly more traffic, developers need tools to help with observability beyond examining individual datapoints. LangSmith provides monitoring charts tracking LLM-specific statistics like number of traces, feedback, time-to-first-token, costs, latency, and error rates over time.

For the research agent in production, monitor:

• Throughput: How many research queries per hour?

• Latency: Average time from query to final summary

• Loop frequency: What percentage of runs trigger the “NEED_MORE_PAPERS” loop?

• Cost: Total tokens consumed across all LLM calls

• Error rate: How often does ArXiv search fail or summarization timeout?

• User feedback: Are summaries helpful? (via feedback API)

A/B Testing and Metadata Grouping

Developers can group monitoring charts by any metadata attribute and instantly get grouped charts over time, allowing experimentation with different LLMs or prompts and tracking their performance.

A typical example experiment for our research problem would be: does GPT-4 or Claude produce better research summaries?

Interactive Drill-Down from Dashboards

Developers can hover over a datapoint in the monitoring chart and click it to be led back to the runs table with a filtered view.

Imagine noticing a latency spike at 2 p.m. reaching eight seconds, three times the usual. Clicking that point reveals runs from the same window, sorted by delay. The slowest trace might show a six-second lag in the researcher node caused by a sluggish ArXiv API. This direct path from metric to root cause greatly speeds up debugging.

Zero-Latency Async Tracing

LangSmith does not add any latency to applications. The LangSmith SDK uses a callback handler that sends traces to a trace collector running as an async, distributed process. In LangChain Python, LangSmith’s tracing is done in a background thread to avoid obstructing production applications.

Together, LangGraph Studio and LangSmith form a complete development lifecycle toolkit: Studio for rapid local iteration and visual debugging, LangSmith for production observability and performance optimization. This separation of concerns, development versus deployment, ensures agents can be built quickly and deployed confidently.

Key Takeaways

✓ Graph-based architectures fundamentally transform agent capabilities beyond linear chains: Unlike sequential chains that execute predetermined steps, graph-based systems enable non-linear exploration, conditional branching, iterative refinement, and state-aware decision-making. This is essential for complex workflows where the path forward isn’t predetermined.

✓ State management creates persistent memory that enables true multi-step reasoning: TypedDict-based state schemas allow agents to carry context across nodes, track progress, maintain conversation history, and make informed decisions based on accumulated knowledge rather than treating each step in isolation.

✓ Modular node architecture and conditional edges enable adaptive, self-correcting workflows: Single-responsibility nodes combined with routing functions create transparent reasoning pipelines that can autonomously decide next steps, loop back when needed, and adapt strategies based on intermediate results, all without brittle retry logic.

✓ Production agents require specialized observability tools designed for non-deterministic AI systems: LangGraph Studio provides visual debugging, interactive state inspection, and hot-reloading for rapid development iteration, while LangSmith delivers end-to-end tracing, production monitoring dashboards, A/B testing capabilities, and zero-latency async tracking. Together, they address the unique challenges of multi-step LLM reasoning that traditional software monitoring cannot handle.

References

[1] LangChain Documentation. (2025). LangGraph: Build resilient language agents as graphs.

[2] LangChain Documentation. (2025). LangChain: Building applications with LLMs through composability.

[3] LangSmith Documentation. (2025). LangSmith: Platform for debugging, testing, evaluating, and monitoring LLM applications.

[4] LangGraph Studio Documentation. (2025). LangGraph Studio: A specialized IDE for developing LangGraph applications.

[5] OpenAI API Documentation. (2025). OpenAI Platform.

[6] arXiv API Documentation. (2025). arXiv API User’s Manual.

This article is also available as a podcast! If you’re on the go or just want to absorb the content in audio format, you can listen to the full episode below 👇 The podcast is also available on Spotify and Apple Podcasts.

While traditional systems fail loudly with error codes, AI agents fail silently with confident nonsense.

The same prompt that generates brilliant analysis one moment can produce hallucinated data the next, all while monitoring dashboards show green. This fundamental shift from deterministic software to probabilistic reasoning systems demands a new operational paradigm: AgentOps.

Objective

This article focuses on bridging the gap between AgentOps theory and practice, demonstrating how to build truly observable AI agent systems. It is based on a real-world use case: a multi-agent investment brief generator that autonomously fetches market data, analyzes financials, monitors news, and synthesizes investment recommendations.

After reading this article, you will understand:

1. Why agentic systems require fundamentally different operations than traditional ML models? This includes unique failure modes, anomaly taxonomies, and the shift from deterministic to probabilistic debugging

2. How does the AgentOps operational framework operate? This encompasses monitoring, anomaly detection, root cause analysis, and iterative resolution strategies tailored for autonomous systems

3. How to implement production-grade observability using MLflow? This relies on tracing, prompt management, and evaluation capabilities to instrument a four-agent pipeline with comprehensive cognitive state capture.

Prerequisites
- Basic understanding of LLMs and prompt engineering (familiar with concepts like tool calling)

Tools & Libraries
- MLflow, for tracing, prompt management, and evaluation framework
- OpenAI API as LLM providers (principles apply to any LLM)

You can find the code here on GitHub.

Subscribe now

1. From MLOps to AgentOps: The New Operations Paradigm

1.1 Why Do Agent Systems Require Different Operations?

Traditional machine learning operations (MLOps) are built around a straightforward paradigm: a model receives an input, generates a single prediction, and the process ends. This framework works well for static tasks such as identifying whether an image contains a cat or assessing the sentiment of a product review.

Agentic systems, however, break away from this model entirely. Rather than producing fixed outputs, agents operate through dynamic, iterative reasoning cycles that incorporate planning, tool usage, and memory.

This shift introduces non-deterministic behavior: the same input can yield different outputs due to the probabilistic nature of large language models (LLMs), even when model weights and parameters remain unchanged.

While traditional models behave consistently once deployed, an agent’s behavior unfolds dynamically in real time. It is shaped by the interaction between its LLM “brain”, prompt instructions, evolving memory, and the external tools it selects.

With LLM-powered agent systems, the paradigm moves from deploying a static prediction function to orchestrating a cognitive system, less like debugging a computation, and more like debugging a conversation.

1.2 How Does the AgentOps Lifecycle Work?

To meet the demands of such dynamic behavior, the AgentOps lifecycle extends the traditional DevOps framework to fit the realities of autonomous agents.

While both follow a similar four-stage structure—monitoring → anomaly detection → root cause analysis → resolution [1, 3], each phase must be reimagined to handle the probabilistic, iterative, and tool-augmented reasoning loops that define agent behavior.

Let’s dive deeper into each of these stages to understand how AgentOps adapts them for agentic systems.

1. Monitoring: Capturing the Agent’s Cognitive State

• Traditional MLOps monitors model metrics (accuracy, drift), data quality, and infrastructure (latency, throughput).

• Agent systems require a fundamentally different approach: cognitive state traceability, capturing not just what an agent did, but why it made specific decisions [1, 3]. This means instrumenting the agent’s internal reasoning loop to record:

  • Beliefs and context: The agent’s current understanding of the world state, derived from initial inputs and accumulated observations. In the example of the investment agent, this includes known stock prices, market conditions, portfolio constraints, and task goals essentially its “investment thesis” at any given moment.

  • Memory contents: Both short-term (the current context window with recent conversation history) and long-term memory (documents retrieved from vector databases, with their similarity scores and retrieval timestamps). For the investor agent, this might include retrieved earnings reports, analyst recommendations, or historical price data. This reveals whether the agent had access to correct information when making decisions.

  • Reasoning traces: The explicit chain-of-thought that transforms beliefs into actions. This is the “inner monologue” showing step-by-step logic. In the context of the investor agent, this would materialize as follows: “Given current market conditions, I should first check recent earnings... The earnings report shows revenue decline... Therefore I’ll downgrade my recommendation...” Capturing this requires either prompting techniques (explicitly asking the LLM to show its reasoning) or model introspection (extracting hidden layer activations).

  • Actions and observations: Every tool call with its exact parameters, the raw tool output (not summarized), and the agent’s interpretation of that output. For the investor agent: calling get_stock_price, receiving the price data, then interpreting whether it signals a buying opportunity. This creates an audit trail showing whether failures originated from bad tool selection, incorrect parameter formatting, or misinterpretation of valid results.

  • Plan evolution: How the agent’s high-level strategy changes across execution steps. Did it follow its original plan? When and why did it deviate? For the investor agent, it might start planning to analyze tech stocks, but after discovering high volatility, pivot to defensive sectors instead. Tracking plan modifications helps identify planning anomalies versus valid adaptive behavior.

  This fine-grained telemetry allows for high-fidelity tracing of decision paths and facilitates root cause analysis when behaviors deviate.

2. Anomaly Detection: Identifying Failures Beyond Metrics

• MLOps relies on quantitative thresholds to detect issues (performance degradation, distribution shifts, SLA breaches).

• Agentic systems require semantic-level anomaly detection to identify hallucinated or fabricated outputs, invalid tool selections or tool misuse, gaps in memory recall or context drift, and failures in coordination between multiple agents.

• Detection methods range from white-box techniques (analyzing model internals) to black-box approaches (validating outputs against external ground truth).

3. Root Cause Analysis (RCA): Diagnosing Multi-Layered Failures

MLOps RCA traces failures through data quality issues, model degradation, and infrastructure dependencies.

For agents, failures may arise from three distinct layers:

1. System-level: Infrastructure issues, broken APIs, or latency in external services

2. Model-level: Hallucinations, reasoning limitations, or knowledge gaps

3. Orchestration-level: Poor prompt design, suboptimal task decomposition, or ineffective tool usage

A single anomaly may result from compounding issues across these layers.

4. Resolution: Iterative Fixes Requiring Longitudinal Validation

• In conventional systems, fixes can often be verified immediately through tests or re-deployment.

• Agent systems, by contrast, require multi-turn validation: fixes (e.g., updated prompts, new constraints, tool adjustments) must be tested across multiple interaction cycles.

• Side effects may only emerge over time, especially in multi-agent setups or long-horizon tasks. Each resolution should be treated as an experimental intervention, requiring observation, measurement, and potential rollback if new failure modes arise.

2. Building Observable Agent Systems: Tools and Frameworks

2.1 What Actually Fails in Agent Systems?

To build robust and observable agent systems, we first need to understand what can go wrong. Unlike traditional software, where failures often stem from deterministic bugs or infrastructure issues, agent systems introduce an entirely new class of failures rooted in probabilistic reasoning, dynamic decision-making, and evolving memory states.

Wang et al. [1] offer a useful starting point with a taxonomy that distinguishes between two broad categories of anomalies:

• Intra-agent arising within a single agent

• Inter-agent emerging from coordination across agents or between agents and their environment

This distinction is crucial, as each class of failure calls for distinct detection mechanisms and intervention strategies.

Empirical results show that these failures are not rare edge cases. They are inherent to current systems. Even cutting-edge agent frameworks show task success rates ranging from 33% to 95%, depending on complexity [2]. For example:

• On the SWE-bench coding benchmark: 33.83%–75.20%

• In web navigation tasks: 37.20%–95.60% [2]

2.2 What Should AgentOps Tools Monitor?

Modern observability tools for agent systems must balance two competing demands: capturing sufficient detail to enable debugging while avoiding prohibitive storage and performance overhead. Dong et al. [3] propose a comprehensive monitoring taxonomy based on systematic analysis of 17 production observability platforms.

Their framework identifies seven core observability dimensions essential for production agent deployments, ordered by criticality: tracing, monitoring, guardrails, evaluation, feedback, customization and prompt management

• Tracing enables LLM/agent tracing to capture each agent span, a single operation like one LLM call, tool invocation, or retrieval step. This includes the complete execution chain: retrieval operations, LLM calls, tool invocations, evaluation spans, and user feedback.

  • For the investment agent, this would mean tracking the complete flow: fetching market data → retrieving earnings reports → analyzing financials → synthesizing recommendations.

  • This forms the foundation for understanding agent behavior and debugging issues in production.

• Monitoring provides agent analytics dashboards to monitor diverse level and dimension statistics metrics about agents, along with LLM cost management and tracking to monitor token spend with foundation model providers.

  • For the investment agent, this would track how often it calls expensive financial APIs, how many LLM tokens are used per investment brief, and success rates across different market conditions.

  • These capabilities are essential for operational visibility and cost control.

• Guardrails implement predefined rules and constraints to limit agent actions, ensuring safe and predictable behavior, with fallback and escalation paths that provide safe defaults or redirect cases to human operators in ambiguous or risky scenarios.

  • For the investment agent, guardrails might prevent recommendations without sufficient data sources, or escalate to human analysts when market volatility is extreme.

  • Safety mechanisms are critical for production deployments.

• Evaluation allows testing agents against benchmarks and leaderboards by creating datasets, defining metrics, running evaluations, and tracking results over time.

  • This includes evaluating final responses, single-step evaluation to assess any agent step in isolation (e.g., whether it selects the appropriate tool), and trajectory evaluation to verify whether the agent took the expected path of tool calls to arrive at the final answer.

  • For the investment agent, this means verifying: did it fetch the correct financial data? Did it call the news monitoring tool appropriately? Did the final recommendation align with the analysis?

• Feedback mechanisms collect both explicit feedback (directly prompting users for thumbs up or thumbs down ratings) and implicit feedback (measuring user behavior through time spent on page and click-through rate).

  • For the investment agent, portfolio managers might rate the quality of investment briefs, while implicit signals show which recommendations led to actual trades. These enable continuous improvement loops.

• Customization enables the provision, spawning, and deployment of customizable and scalable autonomous agents.

  • Agent capabilities can be extended through toolkits from marketplace integrations, multiple vector databases to improve retrieval performance, and fine-tuned models for business-specific use cases.

  • The investment agent can be customized with specialized financial data providers (Bloomberg, Reuters), sector-specific analysis tools, or fine-tuned models trained on historical investment reports.

• Prompt management supports versioning to track different prompt versions for A/B testing and optimizing agent performance.

  • Prompt playground with model comparisons allows testing different prompts and models before deployment, while prompt injection detection identifies potential code injection and secret leaks.

  • For the investment agent, teams can A/B test different analysis frameworks (”value investing” vs “growth investing” prompts), compare models for financial reasoning quality, and detect attempts to manipulate recommendations through malicious inputs.

2.3 How Do Current AgentOps Tools Compare Across Key Features?

While the previous taxonomy defines what should be monitored in production agent systems, the natural question follows: to what extent do existing tools actually implement these capabilities? To answer this, Dong et al. [3] evaluated production AgentOps platforms against the seven dimensions they identified in their November 2024 study.

The table below presents our comparative analysis, revealing significant heterogeneity in feature coverage across the current tool landscape.

* Added by the author

Note on tool coverage: The rapid evolution of the AgentOps ecosystem means that notable platforms may emerge or mature after any snapshot analysis. For instance, MLflow 3.0 introduced comprehensive GenAI capabilities in mid-2024, with agent-specific features continuing to evolve through late 2024 releases. The table represents the state of tools as documented in Dong et al.’s systematic mapping study conducted through November 2024.

3. MLflow for AgentOps: A Practical Implementation

Having established both the theoretical requirements for AgentOps and the current state of tool support, let’s now turn to practical implementation. This section demonstrates how MLflow’s capabilities address the operational challenges identified earlier through a real-world multi-agent system.

Initially developed for traditional machine learning experiment tracking, MLflow has matured into a robust AgentOps platform. Its latest release introduces key capabilities that address the operational challenges outlined in Sections 1 and 2:

• Execution Tracing: It captures detailed records of agent actions, tool usage, and decision paths

• Prompt Management: It enables versioning, comparison, and lifecycle control of behavioral instructions

• Evaluation Frameworks: It supports continuous, automated quality assessments through customizable metrics

These features provide the foundation for building, monitoring, and refining complex agentic systems.

The Investment Brief Agent System

To demonstrate these capabilities, consider a multi-agent system designed for generating structured investment briefs. The architecture includes four specialized agents, each responsible for a distinct analytical task:

1. Data Fetch Agent: Retrieves stock prices, financial data, and key performance indicators from external APIs

2. Financial Analysis Agent: Computes valuation metrics, financial health scores, and historical trends

3. News Monitoring Agent: Identifies and summarizes relevant market news and sentiment shifts

4. Investment Reasoning Agent: Synthesizes structured and unstructured inputs into coherent investment recommendations

This architecture exhibits common failure modes: hallucinated prices, API errors, context loss in long reports, and flawed synthesis, making comprehensive observability essential.

You can find the code here on GitHub.

3.1 How To Trace Agent Decision-Making?

The first operational challenge is understanding how agents actually make decisions. Traditional application monitoring falls short here as logging “Service A called Service B” reveals little about why an agent chose a particular tool, what information it considered, or how it reasoned from inputs to conclusions.

MLflow’s tracing addresses this gap by capturing not just what agents do, but why they make specific decisions. Unlike traditional application tracing (Service A → Service B), agent tracing records cognitive flows using hierarchical spans [3, 4]:

Understanding the Span Hierarchy

Think of spans as nested containers that capture different levels of agent cognition, from high-level task execution down to individual model calls:

• Agent span sits at the top, representing the complete task execution. It captures the agent’s role (e.g., “Financial Analyst”), persona (e.g., “conservative value investor”), and overall success metrics.

• Reasoning spans record individual thinking steps. For the Financial Analysis Agent, this might capture: “Given P/E ratio of 23.4 and industry average of 18.2, stock appears overvalued relative to peers.”

• Planning spans document action selection decisions. When the agent decides “I need current news sentiment before finalizing my recommendation,” the span records both the goal (assess sentiment) and constraints (only articles from last 30 days).

• Tool spans track external interactions. When calling a financial data API, the span captures not just the raw API response, but how the agent interpreted it: “Retrieved revenue growth of 15% YoY, indicating strong expansion.”

• LLM spans provide the finest granularity, recording every model inference with complete prompts, generation parameters (temperature, max tokens), and exact token counts for cost attribution.

The figure below illustrates these concepts in practice. When a user asks “What are Amazon’s main revenue streams and which segment generates the highest profit margins?”, MLflow captures the complete cognitive flow:

With MLFlow, enabling comprehensive tracing requires just a single line of code. Consider this simple example:

3.2 How To Manage Prompts as Production Artifacts?

Prompts control agent behavior, defining roles, boundaries, quality standards, and safety constraints. Yet unlike code, they’ve historically lacked version control and systematic testing. MLflow’s Prompt Registry treats prompts as versioned, production artifacts requiring engineering discipline [2, 4].

Version Control and Rollback

The registry stores complete version history with timestamps, authors, and change descriptions [4]. Each version links to performance metrics: success rates, quality scores, costs, and latency. This enables:

• A/B testing: Run multiple prompt versions simultaneously on production traffic, comparing outcomes statistically to identify which performs best

• Instant rollback: When issues arise, revert to previous stable versions immediately (milliseconds, no redeployment)

• Performance tracking: Query “Which prompt achieved highest quality?” or “Most cost-efficient prompt maintaining 90%+ accuracy?”

Testing Gates for Production Promotion

Before any prompt reaches production, it undergoes systematic validation [4]:

• Regression testing: Verify performance against historical test cases to ensure the new version doesn’t reintroduce resolved issues

• Adversarial testing: Challenge the prompt with edge cases—missing data, contradictory information, ambiguous queries—to expose failure modes before they impact users

• Performance profiling: Measure token efficiency and response latency to catch prompts that sacrifice speed or cost-effectiveness for marginal quality gains

Operational Impact

The registry transforms prompt changes from ad-hoc experiments into trackable, testable, reversible interventions with complete audit trails. When failures occur in production, operators can immediately determine whether issues correlate with recent prompt changes versus other factors (model updates, data quality, API failures). Teams naturally adopt code-like discipline: peer reviews for prompt changes, comprehensive documentation of intent and constraints, validation gates before promotion, and staged rollouts (10% of traffic → 50% → 100%)

The following code demonstrates how to register a versioned prompt template for the Financial Analysis Agent:

3.3 How To Evaluate Agent Quality Continuously?

Traditional software either passes or fails tests. In contrast, agent systems operate in a gray area: the same prompt may yield dozens of plausible responses—none strictly “correct,” yet varying widely in quality.

To address this ambiguity, MLflow’s evaluation framework focuses on multi-dimensional quality assessment rather than binary correctness. This enables robust validation even under the non-deterministic behavior typical of LLM-based agents [2, 5].

Evaluation Approaches

MLflow supports three complementary evaluation strategies:

• LLM-as-Judge Metrics

  Large language models act as evaluators, scoring outputs against structured rubrics. This approach captures nuanced attributes such as:

  • Answer correctness (factual accuracy)

  • Semantic similarity (closeness to reference)

  • Context relevance (appropriateness given prompt context)

    These judges approximate human evaluation and support fine-grained quality analysis.

• Heuristic Metrics

  These are rule-based, function-driven scores including Rouge/BLEU for text similarity, Flesch-Kincaid for readability, toxicity for safety, and latency/token_count for performance.

• Custom Scorers

  Developers can define tailored functions to validate tool usage patterns, constraint satisfaction, domain-specific correctness (e.g. financial ratios, legal disclaimers)

Tracking & Iteration

MLflow integrates evaluation tightly into experiment tracking:

• Run-level Logging: All evaluation results are logged per run, enabling direct comparison across models, prompts, or agent strategies.

• Aggregate Agent Scores: System-level dashboards summarize quality, cost, and latency using metrics like token_count, response_time, and score_distribution.

Evaluation-Driven Development

MLflow promotes a rigorous, feedback-driven development loop:

1. Establish baseline metrics on faithfulness, relevance, reasoning quality, etc.

2. Identify failure patterns via traces and evaluation scores.

3. Implement improvements, e.g. new prompts, tools, memory strategies.

4. Re-evaluate to quantify impact and regressions.

Traces capture multi-step behavior—including tool use, intermediate outputs, and decisions—enabling root cause analysis beyond final responses.

The following code demonstrates how few lines of Python enable systematic quality validation with both correctness checking and custom safety criteria:

Key Takeaways

✓ AgentOps represents a fundamental paradigm shift from MLOps: Unlike traditional ML systems that produce single predictions, agent systems require monitoring cognitive state, capturing not just what agents did, but why they made specific decisions through reasoning traces, memory contents, and plan evolution.

✓ Agent failures are inherently different and more complex: While traditional systems fail deterministically with error codes, agents fail silently with confident hallucinations, tool misuse, and context drift. Current production systems show success rates ranging from only 33% to 95%, making robust observability critical.

✓ The AgentOps lifecycle extends DevOps to handle probabilistic reasoning: The four-stage framework (monitoring → anomaly detection → root cause analysis → resolution) must account for semantic-level failures across three distinct layers: system-level (infrastructure), model-level (hallucinations), and orchestration-level (prompt design).

✓ Seven core observability dimensions are essential for production agents: Tracing captures execution flows, monitoring provides operational visibility, guardrails ensure safety, evaluation validates quality, feedback enables improvement loops, customization extends capabilities, and prompt management treats prompts as versioned production artifacts.

✓ MLflow provides comprehensive AgentOps capabilities with minimal code: Hierarchical span tracing captures complete cognitive flows, the Prompt Registry enables version control with A/B testing and instant rollback, and the evaluation framework supports multi-dimensional quality assessment. This enables production-grade observability with just a few lines of Python.

References

[1] Wang, Z., Li, J., Zhou, Q., Si, H., Liu, Y., Li, J., Xie, G., Sun, F., Pei, D., & Pei, C. (2025). A Survey on AgentOps: Categorization, Challenges, and Future Directions. arXiv preprint arXiv:2508.02121.

[2] Kreuzberger, D., Kühl, N., & Hirschl, S. (2023). Machine learning operations (MLOps): Overview, definition, and architecture. IEEE Access, 11, 31866-31879.

[3] Dong, L., Lu, Q., & Zhu, L. (2024). AgentOps: Enabling Observability of LLM Agents. arXiv preprint arXiv:2411.05285.

[4] MLflow Documentation. (2025).

This article is also available as a podcast! If you're on the go or just want to absorb the content in audio format, you can listen to the full episode below 👇 The podcast is also available on Spotify and Apple Podcasts.

Real-world decisions often involve competing objectives. In supply chain management, for instance, increasing sales and minimizing overstock are two goals that frequently conflict. Traditional optimization approaches tend to merge these goals into a single objective by applying fixed weights—an approach that can oversimplify the problem and obscure critical trade-offs.

Multi-objective optimization offers a more nuanced alternative. Rather than combining objectives into a single equation, it treats each one independently. This allows us to better capture the inherent tension between goals and explore a range of optimal trade-offs. The outcome is a Pareto front: a set of solutions where improving one objective necessarily comes at the expense of another. This empowers decision-makers to select the solution that best aligns with their current priorities—without being constrained by arbitrary weightings.

In this article, we show how to approach a two-objective problem using the Python library pymoo, with a focus on the widely adopted and robust Non-dominated Sorting Genetic Algorithm II (NSGA‑II). [1]

The approach is illustrated with a real-world example: optimizing product dispatch across a network of retail locations.

Objective

After reading this article, you will understand:

1. How can a real-world dispatch problem with conflicting objectives be properly modeled and solved?

2. What are the key mechanisms behind NSGA-II, and how does it address the challenge of balancing multiple objectives?

3. How can Multi-Criteria Decision Making (MCDM) methods be used to evaluate the trade-offs between solutions?

This article assumes a basic understanding of optimization concepts (such as objective functions and constraints) and some experience with Python programming.

All experiments were conducted using the Python libraries pymoo and plotly.

You can find the code here on GitHub.

Subscribe now

1. The Problem at Hand: Boosting Sales Without Overstocking

1.1. Why simple dispatch strategies fall short in the real world?

In theory, product dispatch in retail seems to be a straightforward task: distribute products to stores based on forecasted demand. In practice, however, naïve dispatch strategies often fail to deliver efficient and reliable outcomes.

Common baseline

A common baseline strategy would be to send to each store as much as they are forecasted to sell. While it seems reasonable as we avoid overstocking and undersupply, it falls short for the main reasons:

1. It ignores uncertainty: Demand forecasts are never perfect. Some stores will sell more than expected, others less. Relying blindly on forecast values leads to lost sales or unused stock.

Example: A sudden storm hits a region. Stores there sell out of raincoats, while others are overstocked. A rigid dispatch plan would not catch that.

1. It ignores trade-offs: In shortage scenarios, where total available stock is limited, allocating strictly by forecast does not account for priority, sales potential, or stock efficiency.

Example: A flagship store may need to always have products on shelves to protect brand image, even if its forecast is lower than smaller stores.

Single-objective optimization

One might try to overcome this with a single-objective formulation to maximize, such as:

This scalarization reduces the problem to something a solver can optimize, but at the cost of:

• Choosing arbitrary weights (α, β), which strongly bias the result.

• Losing interpretability as it is unclear how much one objective is being sacrificed for the other.

• Finding only one solution, when there are multiple viable trade-offs.

This is a core limitation that traditional methods like linear programming suffer from.

The extra hurdle: dealing with forecast uncertainty

Even with a good model, forecasting is never exact. Errors of ±10–20% are common in retail due to promotions, seasonality, and local events. A truly robust dispatch strategy must account for these uncertainties. This means not by blindly trusting forecasts, but by simulating multiple plausible scenarios and planning for what works on average, or under worst-case assumptions.

1.2. How to account for competing objectives in dispatch problem formulation?

To address the limitations of naïve and single-objective dispatch strategies, we need to acknowledge that retail dispatch is fundamentally a multi-objective problem. There is not a single “best” way to dispatch products but many viable trade-offs between competing goals.

The dispatch problem involves making the following decision: How many units should be sent to each store, given limited overall stock?

Notations

Let’s denote:

• x_i: the number of units dispatched to store i

• d_i: the demand forecast at store

• S: the total number of units available to dispatch

Our two objectives are:

Objectives

For each product and on each day, the two objectives to optimize are:

1. Maximize expected sales

2. Minimize expected overstock

Contraints

Subject to:

These constraints ensure that each store receives a non-negative quantity of products and that the total dispatch does not exceed the available stock

Stochastic optimization

In reality, demand forecasts are not fixed values, but rather estimates subject to error. Each forecast comes with uncertainty due to unpredictable factors like promotions, weather, local events, or competitor actions.

To handle this, we frame the problem as a stochastic optimization problem.

Stochastic optimization refers to optimization problems where some of the inputs are uncertain or random. The goal is to find solutions that perform well on average, or under worst-case scenarios.

Rather than relying on a single point forecast, the forecast is treated as a central estimate around which multiple demand scenarios are generated. These scenarios reflect realistic variations in demand.

One simple way to do this is to assume demand varies around the forecast according to a normal distribution, where the standard deviation is defined using historical forecast errors (e.g. ±15%). This transforms our view of demand from deterministic to probabilistic.

A more advanced and often more accurate approach involves using probabilistic forecasting models, including those based on machine learning. Unlike traditional models that output a single value, these models provide a full distribution or confidence interval, directly incorporating their own uncertainty into the forecast.

With these scenarios in hand whether sampled from a statistical distribution or learned from data, we can evaluate the performance of each candidate dispatch plan. For every scenario, we compute how many units would be sold (based on the realized demand) and how many would be left unsold, indicating overstock. By averaging the results over all scenarios, we obtain expected sales, a more realistic performance metric under uncertainty, and expected overstock, which captures the inefficiency of over-supplying.

This approach supports robust decision-making, ensuring that the dispatch plan performs reliably across a wide range of possible future outcomes, rather than being optimized for a single, potentially inaccurate forecast.

1.3. Looking at the data: forecasts, inventory, and the challenge of allocation

To make this problem concrete, we use a publicly available retail dataset from Kaggle. It simulates a realistic inventory forecasting scenario in a multi-store retail setting and provides the essential building blocks for dispatch optimization.

The dataset includes daily records for multiple products and stores. Here are the key fields we rely on:

• Date: Daily granularity, spanning a consistent time window across all stores and products.

• Store ID and Product ID: Unique identifiers for each store-product pair.

• Inventory Level: The stock available at the beginning of the day in each store.

• Demand Forecast: The predicted demand for the day, derived from past sales trends.

To simplify the problem setup, it is assumed that the stock available for dispatch each day is equal to the total inventory across all stores. In other words, store inventories are treated as if they were pooled into a single centralized warehouse with no lead time. The core objective becomes determining how to redistribute this shared inventory across the stores to best meet daily demand.

To simulate shortage scenarios the total dispatchable stock is randomly reduced on certain days. This models real-world constraints such as supply chain disruptions, delayed replenishments, or capacity limitations. As a result, the optimization must make strategic trade-offs in allocating limited stock across competing demands.

2. Multi-Objective Optimization with NSGA-II: From Pareto Front to Practical Decision

2.1. How does NSGA-II find a diverse set of optimal trade-offs?

In the early 2000s, most multi-objective optimization techniques relied on repeatedly solving single-objective formulations with different weightings to explore trade-offs. This approach was inefficient, highly sensitive to weight selection, and rarely resulted in a well-distributed set of Pareto-optimal solutions.

Recognizing these limitations, Deb et al. introduced NSGA-II (Non-dominated Sorting Genetic Algorithm II) in 2002 [1]. It quickly became one of the most widely used evolutionary algorithms for multi-objective optimization. It remains today a reference due to its speed, simplicity, and ability to discover well-spread solutions in a single run.

NSGA-II: An evolutionary algorithm for multi-objective optimization

NSGA-II belongs to the evolutionary algorithm (EA) family. Like all EAs, it evolves a population of candidate solutions over successive generations, using principles inspired by natural selection. NSGA-II is specifically designed for multi-objective optimization, where the goal is not to find a single best solution, but a set of non-dominated solutions.

In a multi-objective problem, a solution is said to dominate another if it is no worse in all objectives and strictly better in at least one.

Example: A warehouse has 200 units of a product.

✓ Solution A sends 120 to Store A (forecasted demand: 130) and 80 to Store B (forecast: 60). This leads to higher sales but more overstock.

✓ Solution B sends 100 to each. This leads to less overstock and slightly lower expected sales.

✓ Neither is strictly better, both are non-dominated.

Each generation in NSGA-II follows three main steps, adapted for multi-objective problems:

• Selection: NSGA-II selects parents using binary tournament selection, a simple yet effective method where two candidates are randomly chosen, and the more promising one is picked to reproduce. This mimics natural competition in evolutionary processes.

If Solution A leads to better sales with similar overstock than Solution B, it's selected for the next generation.

• Crossover: Selected parents undergo crossover (e.g., simulated binary crossover) to generate “offspring,” combining characteristics of both parents.

Two dispatch plans [120, 80] and [100, 100] may produce [110, 90].

• Mutation: Offspring are randomly mutated to maintain diversity and help explore new areas of the solution space.

[110, 90] mutates to [115, 85], shifting 5 units to better match changing forecasts.

The resulting set of non-dominated solutions forms an approximation of the Pareto front, a curve (or surface) representing the best achievable compromises between conflicting objectives, such as increasing sales while reducing overstock.

The three pillars of NSGA-II: Sorting, crowding, elitism

NSGA-II introduced three important innovations over earlier multi-objective algorithms:

1. Fast non-dominated sorting

One of NSGA-II’s key innovations is a fast and hierarchical sorting algorithm to classify solutions based on Pareto dominance.NSGA-II uses this principle to organize the population into layers, or non-dominated fronts:

• Front 1 contains solutions not dominated by any other — these represent the current approximation of the Pareto front.

• Front 2 contains solutions dominated only by those in Front 1.

• Front 3 is dominated only by solutions in the first two fronts, and so on.

This forms a hierarchical structure: each front reflects a decreasing level of dominance quality. The algorithm uses this structure to assign ranks to each solution, with lower ranks being preferred during selection.

In the original NSGA algorithm (1994), non-dominated sorting had a computational complexity of O(MN^3), which became a bottleneck for large populations. NSGA-II introduces an optimized method that reduces this to O(MN^2), where:

• N is the number of individuals in the population

• M is the number of objectives

This improvement is achieved by tracking the number of solutions that dominate each individual, and maintaining a list of individuals that each solution dominates

The sorting proceeds by iteratively building each front from solutions with zero domination count (i.e., currently non-dominated) and updating the domination counts of others. This approach avoids redundant comparisons and streamlines the entire process.

In simpler terms: instead of exhaustively comparing each solution with every other, NSGA-II efficiently builds up layers of dominance.

This hierarchical and efficient approach makes NSGA-II scalable and practical for solving real-world multi-objective problems with large search spaces, such as our dispatch optimization use case.

2. Crowding Distance

To maintain diversity, NSGA-II does not just select the best-ranked solutions, it also favors those that are less crowded in objective space. This ensures that the algorithm not only converges toward optimal solutions but also explores a wide range of trade-offs.

To do this, NSGA-II introduces the concept of crowding distance, a measure of how isolated a solution is compared to its neighbors within the same Pareto front. For each solution, the algorithm estimates the "space" around it by computing the normalized distance between its nearest neighbors along each objective.

Solutions with higher crowding distance are preferred during selection when multiple individuals have the same dominance rank. This pushes the algorithm to spread solutions across the Pareto front, avoiding premature convergence to any particular region.

Imagine three candidate dispatch strategies:

A: 95% sales, 20% overstock

B: 96% sales, 21% overstock

C: 92% sales, 15% overstock

If A and B are very close in objective space, and C is farther apart, NSGA-II will likely prioritize C even if it’s slightly less optimal, because it provides a different perspective on the trade-off, and helps construct a more representative Pareto front.

Legend:

• f₁, f₂: Objective functions multi-optimization problem.

• Black circles: Non-dominated solutions on the current Pareto front.

• White circles: Dominated or non-selected solutions (not on the current front).

• i-1, i, i+1: Consecutive individuals on the same front, sorted by objective values.

• Cuboid: Represents the crowding region for individual i, bounded by its two neighbors (i-1 and i+1). The crowding distance is computed as the perimeter (or volume in higher dimensions) of this cuboid.

This crowding-based diversity mechanism eliminates the need for manually tuned niche parameters used in earlier algorithms (like the original NSGA), making NSGA-II both simpler to configure and more robust in practice.

• Pₜ: Parent population at generation t.

• Qₜ: Offspring population generated from Pₜ using variation operators (e.g., crossover and mutation).

• Rₜ: Combined population (Pₜ ∪ Qₜ), which is the input for non-dominated sorting.

• F₁, F₂, F₃...: Fronts generated through non-dominated sorting. F₁ contains the best non-dominated solutions, followed by F₂, and so on.

• Crowding Distance Sorting: Within each front, solutions are ranked based on crowding distance to maintain diversity.

• Rejected: Solutions that are not selected for the next generation due to rank or low crowding distance.

• Pₜ₊₁: Next generation population selected from the best-ranked and most diverse solutions.

3. Elitism

NSGA-II introduces elitism to ensure that the best solutions discovered so far are not lost in future generations, a key improvement over earlier algorithms like the original NSGA.

At each generation, the algorithm:

1. Merges parent and offspring populations,

2. Sorts the combined population using non-dominated sorting and crowding distance,

3. Selects the top individuals to carry forward.

This guarantees that high-quality solutions remain in the population, even if new offspring don’t outperform them, helping the algorithm converge more reliably toward the Pareto front.

In a dispatch context, elitism ensures that a strong dispatch plan — say, one that achieves excellent sales with acceptable overstock — isn’t accidentally discarded as the algorithm evolves.

2.2. How can Multi-Criteria Decision Making (MCDM) techniques guide final selection?

NSGA-II returns a set of non-dominated solutions that represent trade-offs between objectives (e.g., maximizing sales vs. minimizing overstock).

However, in practice, decision-makers need to select a single solution to implement. This is where Multi-Criteria Decision Making (MCDM) techniques come in.

Two commonly used methods, both supported in pymoo, are the Augmented Scalarization Function (ASF) and the Pseudo-Weights approach. Each offers a different way to apply preferences and identify the most suitable dispatch plan.

Augmented Scalarization Function (ASF)

ASF selects the solution closest to the ideal point in objective space, based on a weight vector w = (w_1, w_2) that reflects the relative importance of each objective.

The ASF for a solution x with objective values f(x) is defined as:

Where:

• f_i(x): objective i value of solution x

• z*: the ideal (best) value for each objective across the front

• w_i: weight assigned to objective i

This approach helps select a balanced solution that minimizes the worst deviation from the ideal.

Pseudo-Weights

Pseudo-weights, unlike ASF weights, are not specified by the user. Instead, they are derived from each solution itself, and represent the relative importance of each objective based on the solution’s position within the Pareto front.

For a solution x, with the objectives f_i(x), the pseudo-weight w_i(x) is:

Where:

• f_i(max), f_i(min): max and min values of objective i in the Pareto set

• M: is the total number of objectives.

This yields a normalized weight vector for each solution, summing to 1, which reflects the solution’s implicit prioritization of the objectives, that is, which objective it has improved more relative to the others.

To use pseudo-weights in decision-making, we compute the pseudo-weight vector for each solution and compare it to a user-defined preference vector (e.g., 70% sales, 30% overstock). The solution whose pseudo-weight vector is most aligned with this preference — typically via cosine similarity or angle minimization is selected.

In essence, pseudo-weights allow you to find the solution on the Pareto front that already reflects your desired trade-off, without relying on geometric proximity to the ideal point as ASF does.

Examples

Imagine three dispatch plans:

• Solution A: 90 sales, 8 overstock

• Solution B: 88 sales, 5 overstock

• Solution C: 85 sales, 3 overstock

Suppose your priority is to maximize sales while still keeping overstock under control — say, 80% importance on sales, 20% on overstock.

• ASF will select the solution closest to the ideal ([max sales = 90, min overstock = 3]), considering your weights. In this case, Solution A will likely be chosen, as it achieves the highest sales and its overstock is within a tolerable range.

• Pseudo-weights, however, analyze each solution’s relative improvement across objectives. Solution C reduces overstock significantly compared to the others, but at the cost of substantial sales loss. Solution B offers a more balanced trade-off; a moderate drop in sales for a meaningful drop in overstock.

In this case, pseudo-weights may favor Solution B, as it better matches the direction of your desired trade-off, even if it’s not closest to the ideal.

3. Practical Example: Dispatch Planning Using NSGA-II and MCDM Methods

3.1. How is NSGA-II applied to model the dispatch challenge?

To define the dispatch problem, the DispatchProblem class is introduced to encapsulate the dispatch optimization logic. It inherits from pymoo's ElementwiseProblem, which allows evaluating each solution (i.e., a dispatch plan across all stores) individually.

Upon initialization, we specify key attributes:

• The number of decision variables (n_var, one per store)

• The number of objectives (n_obj = 2)

• One constraint (n_constr = 1) ensuring that total dispatch does not exceed available stock

• Lower (xl) and upper (xu) bounds for the decision variables

The core of the class is the _evaluate method, which is called by pymoo for each candidate solution. In the code, the evaluation is delegated to either a deterministic or a stochastic function depending on whether demand uncertainty is enabled.

The logic is split into two methods depending on whether demand uncertainty is enabled:

• _evaluate_simple: used in the deterministic case, where forecasts are assumed exact.

• _evaluate_with_demand_uncertainty: used when uncertainty is modeled by simulating multiple demand scenarios and evaluating expected outcomes.

3.2. What does a typical optimized dispatch plan look like?

Once the dispatch problem is formulated, we use NSGA-II to explore and optimize possible dispatch plans under uncertainty.

We begin by defining the core components of the evolutionary search:

• FloatRandomSampling() initializes a population of candidate dispatch plans with continuous random values.

• SBX (Simulated Binary Crossover) combines pairs of parent solutions to generate diverse offspring.

• PM (Polynomial Mutation) introduces subtle random variations to promote exploration and prevent premature convergence.

These operators are passed into the NSGA2 algorithm, which evolves a population of dispatch strategies over several generations:

This configuration enables us to explore a wide range of dispatch strategies and ensures that redundant solutions are filtered out to maintain diversity across generations.

We then define the actual dispatch optimization problem. Here, demand is uncertain and evaluated through stochastic simulation using multiple scenarios.

This setup allows the optimizer to evaluate the robustness of a dispatch plan under uncertainty. Demand forecasts are assumed uncertain and simulated using 100 scenarios with 35% relative error.

Finally, pymoo's minimize() is used to run the optimization over 150 generations:

pymoo automatically tracks a range of performance indicators that can be visualized by setting verbose=True

• n_gen – generation number

• n_eval – number of function evaluations

• n_nds – number of non-dominated (Pareto optimal) solutions

• cv_min / cv_avg – minimum and average constraint violations

• eps – epsilon-progress indicator (convergence metric)

• indicator: ideal, nadir, and f are metrics describing the objective space

These logs help you evaluate how well the algorithm is converging, whether the diversity is preserved, and how many valid trade-offs are found.

💡 Tip: If the Pareto front stops improving or the log values plateau early, it may be a sign that your evolutionary settings (population size, mutation rate, number of scenarios) need to be adjusted.

The table below shows the dispatch results for January 4th. The optimizer allocates the available 1154 units across stores in close alignment with forecasted demand, balancing sales performance with overstock risk. The sales figures confirm that the allocation respects both demand and capacity constraints.

3.3. How are trade-offs prioritized using ASF and Pseudo-Weights?

In this first graph, we observe the full Pareto front returned by NSGA-II, plotted in the original objective space. The x-axis represents negative expected sales (since pymoo minimizes objectives), and the y-axis shows expected overstock in units. Each blue dot corresponds to a non-dominated solution.

• The green star marks the ideal point, representing perfect sales with zero overstock.

• The red pentagon indicates the nadir point, the worst observed values on the front.

Figure 5 illustrates the range of trade-offs discovered by the algorithm. As we move from left to right, we gain sales but incur more overstock, a classic optimization tension. It also gives a sense of how much room there is to maneuver between the two objectives.

Figure 6 builds on the first, highlighting two selected solutions:

• The red cross shows the solution chosen using the Augmented Scalarization Function (ASF), which minimizes the weighted distance to the ideal point. The weights assigned to the objectives are 70% for sales and 30% for overstock.

• The orange cross shows the one selected using pseudo-weights, which aims to match the decision-maker’s preferred trade-off direction. The same weights are used.

Here, we see that ASF selects a solution closer to the ideal in terms of raw distance, while the pseudo-weights approach picks a solution that sacrifices a bit more on overstock to preserve sales, offering a directional balance that might better reflect certain business goals.

This graph emphasizes how both methods provide valid but distinct interpretations of what makes a solution “best.”

Finally, Figure 7 shows the same ASF and pseudo-weight solutions, but in a scaled objective space (both axes normalized between 0 and 1). This view helps remove bias due to differing units or magnitudes in the objectives.

Key Takeaways

✓ Multi-objective optimization provides a powerful framework for addressing real-world trade-offs in dispatch planning, especially when goals like maximizing sales and minimizing overstock conflict.

✓ NSGA-II is well-suited for retail scenarios thanks to its ability to generate a diverse set of Pareto-optimal solutions in a single run, without needing to collapse objectives into a single weighted score.

✓ The implementation in pymoo makes it easy to handle both deterministic and uncertain demand through custom evaluation logic, making the approach robust to real-world variability.

✓ When it comes to selecting a final dispatch plan, Augmented Scalarization Function (ASF) and Pseudo-Weights offer two distinct perspectives: ASF favors proximity to the ideal point, while Pseudo-Weights emphasize solutions that match a preferred trade-off direction.

✓ Visualizing the Pareto front, both in raw and normalized form, is key to interpreting results and communicating choices to stakeholders. These plots make the trade-offs and selection rationale transparent and defensible.

References

[1] K. Deb, A. Pratap, S. Agarwal and T. Meyarivan, "A fast and elitist multiobjective genetic algorithm: NSGA-II" in IEEE Transactions on Evolutionary Computation

[2] Kaggle, Retail Store Inventory Forecasting Dataset, A synthetic dataset for practicing inventory management and demand forecasting.

[3] Pymoo documentation, NSGA-II: Non-dominated Sorting Genetic Algorithm

Foundation models have revolutionized NLP by handling tasks like translation, summarization, and reasoning with a single architecture and no task-specific retraining. Now, this paradigm is reaching time series forecasting, where models like TimeGPT, developed by Nixtla, are pushing beyond traditional statistical methods with architectures tailored for temporal data.

Unlike classic forecasting models which rely on domain-specific assumptions or handcrafted features, TimeGPT leverages the Transformer architecture to learn general temporal dynamics directly from large-scale time series corpora. Its strength lies in zero-shot or few-shot performance: it requires little to no historical context to generate accurate forecasts, even in unseen domains.

This article dives into TimeGPT, its design philosophy, forecasting abilities, and practical use in real-world scenarios like sales prediction. It explores its ease of use, performance, and how it compares to both classical and LLM-based approaches such as TimeLLM.

Objective

After reading this article, you will understand:

1. What makes TimeGPT different from other forecasting models?

2. How to use TimeGPT in a zero-shot or few-shot setting for daily time series data?

3. How it performs on the M5 dataset compared to traditional and deep learning models?

This article is the second in a series on advanced time series modeling. If you haven’t already, check out the first part on TimeLLM, Beyond Text: Leveraging LLMs for Time Series Forecasting (Part 1/2).

Prerequisites: No prior experience with LLMs or time series forecasting is required. Some familiarity with Python and pandas will help you follow the code examples.

Tools & libraries: The code used relies on nixtla ecosystem for accessing TimeGPT, pandas for data handling, and plotly for visualization. The dataset used is the M5 competition dataset featuring Walmart daily sales.

You can find the code here on GitHub.

Subscribe now

1. TimeGPT Explained: Foundation Models for Zero-Shot Time Series Forecasting

What makes TimeGPT different?

TimeGPT is the first foundation model built specifically for time series forecasting. It delivers zero-shot predictions, meaning it can accurately forecast entirely new, unseen time series without any retraining or task-specific tuning. This contrasts sharply with traditional models, which typically require manual optimization and feature engineering for each new dataset. [1]

Key differentiators:

• Zero-shot forecasting: TimeGPT can forecast without retraining or fine-tuning.

• Cross-domain generalization: It performs well across diverse areas such as finance, energy, healthcare, and IoT.

• Operational simplicity: Once the data is properly formatted, forecasting is reduced to a single call. No model selection, hyperparameter search, or manual tuning required.

How is TimeGPT built and trained?

TimeGPT is not based on any existing large language model (LLM). But, it follows the same core training principle, which involves scaling up transformer architectures on large datasets, but it is specifically designed and trained for numerical time series data. [1]

Model architecture

TimeGPT is based on a Transformer encoder-decoder architecture adapted for time series forecasting. It is designed to handle time series of different frequencies, lengths, and structures, including varying seasonality, sparsity, and noise. It processes historical values and exogenous covariates, using positional encoding to retain temporal order. The model relies on multi-head attention, residual connections, and layer normalization to capture both short- and long-term dependencies. CNN layers are also included to extract local temporal patterns. A final linear layer projects the output to the forecast horizon, producing predictions along with uncertainty intervals. [1]

However, the original article lacks key details about the model’s design, such as the number of layers, attention heads, and input window size. It also doesn’t specify how exogenous variables are integrated or whether the model uses causal masking and sequence-to-sequence or autoregressive decoding.

Training Dataset

TimeGPT was trained on over 100 billion data points from the largest publicly available time series dataset to date [1], spanning domains like finance, healthcare, IoT, transport, weather, web traffic, and energy. The dataset includes mixed frequencies (hourly to monthly), varying noise levels, outliers, trends, and seasonality, and heterogeneous patterns across domains

Preprocessing was minimal, focusing on formatting and handling missing values while preserving raw signal variability to enhance robustness.

TimeGPT was trained in PyTorch over a few days on NVIDIA A10G GPUs using the Adam optimizer with learning rate decay and large batch sizes. [1] In contrast, models like GPT-3 are trained for weeks on thousands of GPUs with more complex setups. While both use Transformers, TimeGPT remains lightweight and specialized for time series compared large, language-focused LLMs.

How well does TimeGPT perform in practice?

TimeGPT was evaluated in a true zero-shot setting—no fine-tuning, just direct inference on over 300,000 unseen time series. [1]

Performance was measured using two normalized metrics: Relative Mean Absolute Error (rMAE) and Relative Root Mean Squared Error (rRMSE). These scores are relative in the sense that they are normalized against the error of a seasonal naive baseline, which assumes each future value repeats the last observed value from the same season (e.g., predicting next January based on the previous January). A score of 1.0 indicates parity with this naive benchmark, while values below 1.0 reflect superior performance. This normalization enables consistent comparisons across datasets with different scales, units, and temporal patterns.

Results

• TimeGPT ranks 1st in 4 out of 8 metrics (rMAE and rRMSE across monthly and weekly frequencies), and is among the top 3 in all others. It demonstrates strong and consistent zero-shot performance across time scales.

• Monthly and Weekly Forecasts: TimeGPT clearly outperforms all other models, including specialized statistical (e.g. DOTheta) and deep learning methods (e.g. NHITS, TFT).

• Daily and Hourly Forecasts: While NHITS and LGBM are better on some hourly metrics, TimeGPT remains competitive and outperforms most traditional models.

Overall, TimeGPT shows that a single, general-purpose model can match or surpass domain-tuned baselines, even without fine-tuning, validating its foundation model approach.

Note:

• TimeGPT delivers rapid inference at just 0.6 milliseconds per time series. This is faster than most statistical models, even when parallelized.

• TimeGPT can also be fine-tuned on specific domains to further improve accuracy.

• It supports uncertainty quantification through conformal prediction, enabling reliable prediction intervals without relying on strong distributional assumptions.

2. Foundation Models in Action: Forecasting with TimeGPT and Nixtla

Use case: M5 store-level forecasting with TimeLLM and benchmarks

This use case follows the same setup as described in the article Beyond Text: Leveraging LLMs for Time Series Forecasting (Part 1/2). Please refer to the article for a full description of the dataset, benchmark models, and experimental protocol.

Problem Statement

The task is to forecast daily store-level turnover over a 14-day horizon using historical sales data and external features such as calendar events from M5 dataset [14]. Each store is modeled independently for the target variable, while shared temporal patterns are leveraged across series.

Learning and Evaluation Framework

The last 14 days of each store’s time series are held out as a test set. Models are trained on historical data only. Evaluation is performed using the Mean Absolute Percentage Error (MAPE).

Evaluation strategy

To assess the value of TimeGPT's architecture and feature integration, three progressively enriched configurations were tested:

1. TimeGPT (baseline): The base timegpt-1 model trained only on historical sales, with no additional features.

2. TimeGPT (with date features): The same model (timegpt-1) enriched with automatically extracted calendar features (e.g., day of week, holiday flags).

3. TimeGPT (long horizon): The timegpt-1-long-horizon variant, which uses a model architecture tailored for longer seasonal dependencies. Feature setup remains identical to the previous step.

4. TimeGPT (with exogenous features): This final configuration combines the long-horizon model with both calendar features and external event indicators (e.g., bank holiday flags) from the M5 dataset.

Benchmark Models

• Statistical models: AutoARIMA, AutoETS, AutoCES, MSTL, Theta

• Deep learning models: NHITS, NBEATS (trained with quantile loss and enriched with dynamic exogenous features)

How is TimeGPT used in practice with Nixtla libraries?

Prerequisite:

Running TimeGPT requires access to the Nixtla Python client, along with an API key, which can be obtained by registering on nixtla.io.

The input data must include:

• A timestamp column (time_col, e.g. "ds")

• A target column (target_col, e.g. "y")

• An identifier column for multiseries (id_col, e.g. "unique_id")

• Optionally, any exogenous variables for both historical and future periods

Example: forecasting with TimeGPT

Here’s the code used in this benchmark to forecast a 14-day horizon using the timegpt-1-long-horizon model, fine-tuned on recent data. For full code and implementation details, see the GitHub repository.

Key parameters explained

• df / X_df: Historical and future data, including exogenous variables if applicable.

• h: Forecast horizon (e.g., 14 for 14-day prediction).

• finetune_steps and finetune_depth: Enable model adaptation to recent trends; depth 3 is a moderate setting.

• finetune_loss: Loss used during finetuning (e.g., "rmse").

• model: Use "timegpt-1-long-horizon" for multi-seasonal or longer horizons.

• date_features: Automatically extract calendar features (e.g., weekday, month).

• level: Confidence interval for probabilistic forecasting (e.g., 90%).

For a complete reference of all parameters, see the official TimeGPT client documentation.

How does TimeLLM perform compared to benchmark models?

TimeGPT outperforms both statistical and deep learning models on the M5 store-level forecasting task. The best configuration, combining a long-horizon architecture, calendar features, and external event indicators, achieves a MAPE of 6.2%, outperforming traditional baselines like AutoARIMA (9.2%) and advanced neural models like NHITS (7.8%).

Mean MAPE comparison

TimeGPT shows consistent improvements as feature complexity increases. The full configuration with exogenous features achieves the lowest average MAPE, while even the baseline version remains competitive with other modern models.

MAPE distribution across models

MAPE distributions reveal that TimeGPT not only improves on average performance but also reduces error variability across stores. The most robust results are obtained with the full-featured model, while statistical methods exhibit wider and less consistent error ranges.

Forecast example with confidence interval

An example forecast for store WI_1 illustrates how TimeGPT captures weekly sales dynamics and aligns closely with the observed test values. The 90% prediction interval reflects the model’s calibrated uncertainty over the forecast window.

How interpretable is TimeGPT?

While accuracy is key, explainability is critical for real-world adoption. TimeGPT delivers SHAP-based feature attributions, revealing which variables drive predictions, both overall and at each time step.

Extracting SHAP values

By enabling the feature_contributions argument, TimeGPT returns SHAP values alongside forecasts. These quantify how much each feature contributes to the predicted value compared to a base reference.

The SHAP values can then be accessed using:

Global importance of features

To understand which features generally drive model predictions, a bar plot of mean absolute SHAP values can be generated:

This bar chart ranks features by their average impact on the forecast. Features related to calendar structure and special events emerge as dominant drivers of sales volume across stores and dates.

Local interpretation of a forecast

To dive deeper, a waterfall plot can explain a single forecast by showing the additive effect of each input feature on the final predicted value.

The waterfall plot breaks down how each feature shifts the model output from the base value to the final forecast. Positive and negative contributions are clearly distinguished, offering transparent reasoning behind individual predictions.

Key Takeaways

✓ TimeGPT brings foundation model capabilities to time series forecasting

It enables accurate zero-shot and few-shot predictions across domains without the need for retraining or manual tuning, marking a shift from traditional, task-specific forecasting pipelines.

✓ Architecture and feature integration drive consistent gains

TimeGPT's performance improves systematically with richer temporal structures (e.g., calendar features) and external signals, while maintaining simplicity of use through a single API call.

✓ Built-in explainability enhances trust and usability

With native SHAP integration, TimeGPT not only predicts but also explains. This enables the interpretation of the influence of features on each forecast at both global and individual levels.

References

[1] A. Garza, C. Challu, M. Mergenthaler-Canseco (2023), TimeGPT-1: Foundation Model for Time Series Forecasting, arXiv

[2] Wenlong Liao, Fernando Porte-Agel, Jiannong Fang, Christian Rehtanz, Shouxiang Wang, Dechang Yang, Zhe Yang, *TimeGPT in Load Forecasting: A Large Time Series Model Perspective,* arXiv

[3] Congxi Xiao, Jingbo Zhou, Yixiong Xiao, Xinjiang Lu, Le Zhang, Hui Xiong, TimeFound: A Foundation Model for Time Series Forecasting, arXiv

[4] Hua Tang, Chong Zhang, Mingyu Jin, Qinkai Yu, Zhenting Wang, Xiaobo Jin, Yongfeng Zhang, Mengnan Du, Time Series Forecasting with LLMs: Understanding and Enhancing Model Capabilities, arXiv

[5] Jin, M., Wang, S., Ma, L., et al. (2023), Time-LLM: Time Series Forecasting by Reprogramming Large Language Models, arXiv

[6] Xiyuan Zhang, Ranak Roy Chowdhury, Rajesh K. Gupta, Jingbo Shang, Large Language Models for Time Series: A Survey, arXiv

[7] C. Liu, S. Zhou, Q. Xu, H. Miao, C. Long, Z. Li, R. Zhao (2025), Towards Cross-Modality Modeling for Time Series Analytics: A Survey in the LLM Era, arXiv

[8] T. Kim, J. Kim, Y. Tae, C. Park, J. Choi, J. Choo (2022), *Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift,* ICLR 2022, Paper with code

[9] Nixtla, TimeGPT-1 documentation

[10] Ming Jin, Time-LLM: Time Series Forecasting by Reprogramming Large Language Models, MLLM Talk, Youtube

[11] Adrien Nav, *TimeLLM - Time Series Forecasting Model,* Youtube

[12] Marco Peixeiro, Time-LLM: Reprogram an LLM for Time Series Forecasting, Medium

[13] S. Makridakis, E. Spiliotis, V. Assimakopoulos, *M5 accuracy competition: Results, findings, and conclusions,* ScienceDirect

[14] Kaggle, M5 Forecasting

Large Language Models (LLMs) have reshaped the AI landscape by enabling a single architecture to handle a wide range of tasks, often outperforming models built for specific problems. Their strength lies in generalization: trained on massive datasets with simple objectives, LLMs can tackle translation, question answering, reasoning, and more, all without task-specific tuning.

This flexibility has allowed LLMs to expand far beyond natural language processing. They now power applications in code generation, multimodal learning, and even computer vision, where they serve as universal backbones.

A new frontier is emerging: time series forecasting. Traditionally dominated by specialized statistical models, forecasting is now being reimagined through foundation models trained on sequential patterns.

Objective

This article explores how Large Language Models (LLMs) can be applied to time series forecasting, with a particular focus on TimeLLM, a model that repurposes language modeling techniques for temporal data. It delves into both the theoretical foundations of TimeLLM and its practical implementation using the Nixtla ecosystem. The methodology is illustrated through a concrete use case: sales forecasting on the M5 competition dataset, which features daily Walmart sales data across thousands of items and stores.

After reading this article, you will understand:

• How does TimeLLM repurpose language models to tackle time series forecasting?

• How performant is it compared to other classic approaches?

• How can these models be applied in practice to deliver fast, accurate time series forecasts?

This article is the first in a series on advanced time series modeling. Be sure to explore the others for a complete understanding.

Prerequisites: The article is designed to be accessible to readers without prior expertise in the subject.

Tools & libraries: The code uses the nixtla library for TimeLLM implementation and Plotly for interactive visualizations. Additional tools include pandas and scikit-learn for data handling and evaluation.

You can find the code here on GitHub.

Subscribe now

1. TimeLLM Explained: Repurposing Language Models for Temporal Forecasting

What is the core idea behind TimeLLM?

LLMs are no longer confined to natural language. In fields like computer vision and code generation, they now serve as powerful general-purpose backbones, outperforming task-specific models. This cross-domain success raises a timely question for time series forecasting: can the pattern recognition and reasoning abilities of LLMs be harnessed for temporal data?

Time series and language share a sequential nature, but differ in structure and encoding. The key challenge is aligning these modalities to allow LLMs to process time series effectively, without modifying the models themselves.

This is the core idea behind TimeLLM, introduced in “Time-LLM: Time Series Forecasting by Reprogramming Language Models” [5]. Rather than redesigning architectures, the authors reprogram the input: transforming time series into token sequences that resemble natural language prompts. This approach enables pretrained language models, like GPT-2, to be used directly for forecasting, leveraging their rich priors and generalization capabilities.

How are time series reprogrammed into language-compatible formats?

This process involves the following key steps:

1. Input Embedding: The first step is to convert continuous time series data into compact patch representations. Each time series is first normalized (zero mean, unit variance) using RevIN [8], to handle distribution shifts. The time series is then segmented into patches, small overlapping or non-overlapping windows of consecutive values.

Original series (T = 8): [95, 96, 97, 98, 100, 99, 98, 97]

Patch length (Lp) = 4

Patches:
- Patch 1: [95, 96, 97, 98]
- Patch 2: [100, 99, 98, 97]

Emdedding dimension (dm) = 3

Patches embedded:
- Patch 1: [0.2, 0.3, 0.4]
- Patch 2: [-0.3, -0.2, -0.1]

Each patch is linearly embedded into a vector of dimension dₘ. This results in patch embeddings.

2. Patch Reprogramming: The next step involves aligning the structure of time series embeddings with the language model’s representation space. This is done by transforming each patch, which is a short segment of the series, into a combination of language-like representations. Instead of retraining the model, TimeLLM selects a small set of prototypes from the LLM’s existing vocabulary, words like “up”, “down”, or “steady” that reflect common time series patterns. Each patch is then matched to these prototypes. For example:

   • A decreasing patch might become “short down”

   • A stable patch could be “steady flat”

   

   This transformation uses attention to choose the most relevant prototypes, converting numeric data into something the LLM can naturally process, just like a sentence. By rephrasing time series in this way, TimeLLM enables a frozen language model to perform forecasting without any architectural changes.

Reprogrammed output:
- Patch 1: [0.2, 0.3, 0.4] → “steady up”
- Patch 2: [-0.3, -0.2, -0.1] → “slow down”

3. Prompting: The next step consists of building a prompt that the LLM can process naturally. TimeLLM adopts the Prompt-as-Prefix (PaP) strategy instead of the Patch-as-Prefix approach, enabling better alignment with the model’s pretraining format.

• Patch-as-Prefix: It translates the past values of the time series directly into natural language and appends a textual instruction asking the LLM to forecast. For example:

"Weekly sales were: 234.5, 238.0, 231.2, 229.8. Predict the next 3 values."

• This method faces challenges:

  • LLMs tokenize numbers in inconsistent ways (e.g., 238.0 might be split into ['2', '38', '.', '0'] depending on the tokenizer).

  • It lacks semantic structure and statistical grounding.

  • It’s hard to scale for long input sequences or high-precision forecasts.

• Prompt-as-Prefix (PaP): It avoids these limitations by enriching the prompt with structured metadata and reprogrammed patch embeddings. It follows a format like:

The M5 dataset contains daily sales records for thousands of products sold across Walmart stores in the United States. Each record includes product-level sales, calendar information, and event markers. Below is the information about the input time series:

• [BEGIN DATA]

• [Domain]: Daily sales of an individual product in a Walmart store.

• [Instruction]: Predict the next <H> days of sales given the previous <T> days attached.

• [Statistics]: The input has a minimum of <min_value>, a maximum of <max_value>, and a median of <median_value> The overall trend is <upward or downward>. The top five lags are <lags>. [END DATA]

4. Output Projection:

   We have seen how to structure and deliver temporal information to an LLM using reprogrammed patches and carefully designed prompts. But how is the model’s output transformed into usable forecasts?

   After the prompt and patch embeddings are passed through the frozen LLM, the model generates a sequence of hidden representations. Only the final part, corresponding to the continuation of the sequence, is used for prediction.

   The remaining outputs are then flattened and passed through a linear projection layer to convert them into real-valued forecasts. This final step completes the decoding process, transforming the LLM’s abstract representations back into interpretable, actionable time series predictions.

⚠️ Important note: TimeLLM does not involve fine-tuning the LLM. The language model (e.g., GPT-2) remains entirely frozen during training. Its weights are not updated. Instead, forecasting is enabled by learning lightweight modules that reprogram the input (e.g., patch embeddings and prompts) into a format the LLM can process. This method bridges zero-shot-style forecasting and few-shot learning, leveraging the LLM’s pretrained reasoning capabilities without expensive retraining or parameter updates.

How well does TimeLLM perform in practice compared to state-of-the-art models?

Learning frameworks

The evaluation of Time-LLM was conducted across a comprehensive set of benchmarks and scenarios, including long-term and short-term forecasting, few-shot learning (with 5–10% of training data), and zero-shot cross-domain transfer.

Datasets

The model was tested on widely accepted datasets such as ETT (ETTh1, ETTm1, etc.), Weather, Traffic, Electricity (ECL), ILI, and M5, using standard metrics including Mean Squared Error (MSE) and Mean Absolute Error (MAE).

Models

It was compared against both classical and recent state-of-the-art models, including transformer-based baselines (PatchTST, FEDformer, Informer), as well as recent LLM-based methods (GPT4TS, LLMTime). Importantly, the comparisons followed unified evaluation protocols from prior works to ensure consistency.

Results

While the reported results show strong and consistent performance, especially in zero-shot and few-shot settings, the claim that “Time-LLM consistently outperforms state-of-the-art forecasting methods by large margins” should be interpreted with nuance.

• The margins are indeed substantial in specific low-data regimes, with reported improvements of up to 33% in MSE, and 14.2% average gain in zero-shot scenarios.

• However, many of these gains are context-dependent: they are most pronounced when the training data is scarce, and less dominant under full supervision.

• Additionally, while the LLM backbone is frozen, the reprogramming components are trained, which means the approach is not strictly zero-shot end-to-end.

• Furthermore, model efficiency and scalability remain important trade-offs compared to domain-specific architectures optimized for speed and resource constraints.

Thus, while the results validate the potential of LLMs in forecasting, they do not fully displace task-specific models in all conditions. This is further supported by the experimental results discussed in the following section.

2. LLMs in Action: Implementing TimeLLM with Nixtla

Use case overview: M5 sales forecasting

About data

To compare the performance of TimeLLM, the M5 dataset was used as a representative real-world benchmark [12].

The data was aggregated per store to forecast daily turnover over a 14-day horizon for 10 stores, each with 1,990 observations. This setup simulates a common business case: making accurate short-term forecasts across multiple small time series with limited historical depth.

While the original M5 competition emphasized item-level forecasting with hierarchical reconciliation, several key takeaways are worth noting. The most impactful strategies included shared learning across related series, the integration of exogenous variables such as calendar events, and the strong performance of hybrid architectures. Hybrid architectures blend the structure of statistical methods with the adaptability of neural networks to capture both predictable patterns and complex dynamics. [11]

Problem statement

The task involves predicting the next 14 days of store-level turnover based on historical sales and external event features. Each store's series is treated independently in terms of target values, but shared representations and temporal structures are leveraged where applicable.

Learning and evaluation framework

For evaluation, the last 14 days of each time series were held out as the test set to ensure models only learned from past data. This simple holdout strategy was preferred over cross-validation, which was less practical given the relatively short length of the series.

Model performance is evaluated using the Mean Absolute Percentage Error (MAPE), which provides a scale-independent measure of accuracy

Benchmark models

To contextualize TimeLLM’s performance, it was compared against two classes of models: traditional statistical methods and modern deep learning architectures.

✔ Statistical Models

• AutoARIMA: Automatically selects ARIMA parameters using AIC minimization. It captures linear autoregressive patterns with integrated differencing and moving average components.

• AutoETS: Automatically fits an exponential smoothing model with components for level, trend, and seasonality. Particularly strong on short-term regular patterns.

• AutoCES (Complex Exponential Smoothing): A generalization of ETS (Error, Trend, Seasonality) designed to better handle nonlinearities and complex seasonality by modeling interactions of components multiplicatively or additively.

• MSTL (Multiple Seasonal-Trend decomposition using Loess): An extension of STL that supports multiple seasonal patterns (e.g., weekly and yearly) and smooths the trend nonparametrically. Especially effective on retail and web traffic series.

• Theta: A simple but powerful model that decomposes the series into theta lines and extrapolates using linear and exponential components. It was the winning model of the M3 forecasting competition.

✔ Deep Learning Models

• NHITS: An interpretable deep learning architecture that learns hierarchical residual blocks over varying temporal scales. It focuses on local level structures and performs well on long-horizon forecasting.

• NBEATS: A pure deep feedforward neural network architecture using backward and forward residual stacks. It does not require prior knowledge of seasonality or trend, making it versatile for various domains.

Note: Both deep learning models were trained using a quantile loss to support probabilistic forecasting, with a confidence level of 90%. In addition, they were enriched with dynamic exogenous variables such as holiday and event indicators, which are known to affect store-level turnover.

How can TimeLLM be implemented using Nixtla's ecosystem?

Nixtla provides a PyTorch Lightning-compatible implementation via the neuralforecast library, allowing data scientists to train TimeLLM models just like any other forecasting model.

1. Forecasting Setup

These parameters define how the historical input is structured before entering the LLM.

• h: Forecast horizon (in our case 14 days).

• input_size: Number of past time steps the model looks at.

• patch_len: Length of each time series segment (patch).

• stride: Step size between patches (controls overlap).

2. LLM and Prompt Configuration

Prompting allows better alignment between numerical patterns and the LLM's pretraining distribution.

• llm: Pretrained language model to use (e.g. 'gpt2', 'gpt2-medium').

• d_llm: Dimensionality of the LLM’s hidden states (e.g. 768 for GPT2, 1024 for GPT2-medium).

• prompt_prefix: Optional natural-language prompt describing the forecasting task and data context.

3. Reprogramming Architecture

These define how numeric input is transformed into text-like embeddings for the LLM.

• d_model: Internal dimension for patch-to-prototype embedding.

• d_ff: Hidden dimension for the patch encoder’s feed-forward layers.

• n_heads: Number of attention heads in the cross-attention layer.

• top_k: Number of LLM token embeddings to use as prototypes.

4. Training Parameters

These parameters control the learning process of the lightweight components surrounding the LLM.

• batch_size, valid_batch_size, windows_batch_size: Batch sizes for training and validation.

• learning_rate: Learning rate for optimizing reprogramming layers (default: 1e-4).

• max_steps: Maximum number of training iterations (e.g. 500–1000).

• dropout: Dropout probability for regularization.

• loss: Loss function (e.g. MAE(), MQLoss(level=[90])).

• val_check_steps: How often to evaluate on the validation set.

6. Additional Controls

These offer extra flexibility and stability in training TimeLLM effectively.

• scaler_type: Input normalization method ('robust', 'minmax', etc.).

• early_stop_patience_steps: Stop training early if no improvement.

• start_padding_enabled: Allow left-padding for short sequences.

• step_size: Window stride for sampling training windows.

• random_seed: Seed for reproducibility.

For more information, refer to the Nixlta documentation [7].

Once the model is defined, it can be trained and used for forecasting through the NeuralForecast framework. The fit() method trains the model on historical data and calibrates 90% prediction intervals using conformal methods based on 10 past windows. The predict() function then produces both point forecasts and associated uncertainty intervals for the next 14 days.

The code display logs that summarize the structure of the TimeLLM model at training time. The GPT-2 backbone (llm) is kept frozen (eval mode) with 124M non-trainable parameters, while the remaining 54M parameters are trainable and belong to the lightweight reprogramming modules. Most of the trainable parameters are concentrated in the mapping_layer, which transforms the LLM outputs into forecast-ready representations. Modules like patch_embedding, reprogramming_layer, and output_projection handle the transformation of time series patches into embeddings aligned with the LLM space. In total, the model includes 178M parameters and is compact enough (713 MB) to run comfortably on standard hardware.

How well does TimeLLM perform?

Empirical results on store-level forecasting

In the evaluation based on the M5 dataset (aggregated at store level, 14-day horizon), TimeLLM performed significantly worse than both classical and deep learning benchmarks. While NBEATS and NHITS achieved MAPE scores around 7-8%, and statistical models like AutoETS reached 8-9%, TimeLLM’s MAPE remained around 12.8%.

Visually, TimeLLM’s forecasts appeared overly smooth, failing to reflect known seasonal dynamics such as weekend demand surges or event-driven peaks. The model’s output consistently regressed toward the historical average, suggesting it lacked responsiveness to calendar effects, even when these were included as binary flags in the input.

This behavior resembled that of a naïve or exponentially smoothed baseline rather than a model designed for pattern-aware generalization. This outcome stands in sharp contrast to the high performance reported in the original TimeLLM paper.

Understanding the discrepancy with the initial paper

In the initial paper “TIME-LLM: Time Series Forecasting by Reprogramming Large Language Models” [3], the model consistently outperforms deep learning and transformer-based methods across various long-range forecasting tasks. According to Table 12 [3]:

• TimeLLM achieves the lowest SMAPE and MASE on quarterly and yearly data.

• It ranks first in average OWA across benchmarks (0.86), outperforming PatchTST, FEDformer, and DLinear.

However, these gains are observed in datasets with high trend/seasonal signal strength and input lengths ranging from 96 to 512 time steps, often evaluated in few-shot or zero-shot configurations.

Several design omissions likely contributed to the underperformance of TimeLLM in our experimentations:

• Limited model capacity: The experimentations used GPT2-small (117M parameters), whereas the paper relies on large-scale LLMs like LLaMA-7B or GPT-4. Smaller models lack the abstraction capacity to generalize across complex time patterns.

• No structured prompt engineering: The prompt used described calendar effects in plain language but omitted statistical context such as trend slope, min/max, median, or autocorrelation lags. These elements are crucial in the original paper.

• Short input sequences: Input sizes in the original benchmarks ranged from 96 to 512 steps, whereas our setup used only 30–60, limiting the LLM’s ability to detect and reason over seasonal cycles.

• No decomposition or reformatting: The data was passed in raw numerical form without trend or seasonal decomposition, which removes useful inductive bias and prevents the LLM from focusing on meaningful variation.

• No input paraphrasing: As highlighted in “Time Series Forecasting with LLMs: Understanding and Enhancing Model Capabilities” [2], rephrasing inputs into natural language (e.g., “sales dropped from 180 to 140”) significantly improves performance. This step was not applied.

• Mismatch in learning regime: TimeLLM is designed for zero-or few-shot generalization. In our case, the model was trained on full-length data for just 40 steps. This is a misaligned setup for what is fundamentally a pretrained LLM.

As a result, the model underperformed by a wide margin and did not demonstrate the generalization capabilities highlighted in the original paper.

Broader Alignment with the Literature

However, these findings are consistent with the limitations outlined in ”Time Series Forecasting with LLMs: Understanding and Enhancing Model Capabilities” [2] which offer a more critical perspective. In that work, LLM-based methods, whether GPT-4, Gemini, or LLaMA-2, only performed competitively under highly structured conditions.

The authors observed:

• Strong results only on datasets with high seasonal and trend strength (e.g., MonthlyMilk, AirPassengers)

• Poor generalization on short or irregular datasets (e.g., Sunspots, HeartRate), where MAPE often exceeded 300% (cf. Table 3 in the paper [2]).

• Significant gains only when input sequences were reformatted into natural language or when external statistical knowledge was included

Our M5 setup aligns with these weaker cases: moderately seasonal, event-sensitive series with limited context and no prompt rephrasing. The absence of input-to-text transformation or decomposed statistics likely contributed to the lack of generalization. In particular, the model showed no capacity to recover or amplify signal from binary event indicators—an issue both papers identify when external knowledge is omitted.

Key Takeaways

✔️ TimeLLM adapts frozen LLMs for time series forecasting by transforming numerical sequences into language-like prompts, allowing temporal predictions without modifying the core model.

✔️ The model shines in zero-shot and few-shot scenarios, with the original paper reporting up to 33% gains in MSE over state-of-the-art baselines when training data is limited.

✔️ The innovation lies in reprogramming, using patching and prompt engineering to align time series data with the LLM’s input expectations—no fine-tuning required.

✔️ However, TimeLLM is more resource-intensive than classical and even deep learning models. Despite not training the LLM, inference and training of the surrounding architecture still demand considerable compute and memory.

Final thoughts: TimeLLM is a strong choice when data is scarce, flexibility is critical, or you need fast prototyping across domains. For production-grade or resource-constrained environments, leaner task-specific models often remain more efficient and accurate.

References

[1] Congxi Xiao, Jingbo Zhou, Yixiong Xiao, Xinjiang Lu, Le Zhang, Hui Xiong, TimeFound: A Foundation Model for Time Series Forecasting, arXiv

[2] Hua Tang, Chong Zhang, Mingyu Jin, Qinkai Yu, Zhenting Wang, Xiaobo Jin, Yongfeng Zhang, Mengnan Du, Time Series Forecasting with LLMs: Understanding and Enhancing Model Capabilities, arXiv

[3] Jin, M., Wang, S., Ma, L., et al. (2023), Time-LLM: Time Series Forecasting by Reprogramming Large Language Models, arXiv

[4] Xiyuan Zhang, Ranak Roy Chowdhury, Rajesh K. Gupta, Jingbo Shang, Large Language Models for Time Series: A Survey, arXiv

[5] C. Liu, S. Zhou, Q. Xu, H. Miao, C. Long, Z. Li, R. Zhao (2025), Towards Cross-Modality Modeling for Time Series Analytics: A Survey in the LLM Era, arXiv

[6] T. Kim, J. Kim, Y. Tae, C. Park, J. Choi, J. Choo (2022), *Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift,* ICLR 2022, Paper with code

[7] Nixtla, TimeLLM documentation

[8] Ming Jin, Time-LLM: Time Series Forecasting by Reprogramming Large Language Models, MLLM Talk, Youtube

[9] Adrien Nav, *TimeLLM - Time Series Forecasting Model,* Youtube

[10] Marco Peixeiro, Time-LLM: Reprogram an LLM for Time Series Forecasting, Medium

[11] S. Makridakis, E. Spiliotis, V. Assimakopoulos, *M5 accuracy competition: Results, findings, and conclusions,* ScienceDirect

[12] Kaggle, M5 Forecasting

The rise of Large Language Models (LLMs) has revolutionized the way we extract information from text and interact with it. However, despite their impressive capabilities, LLMs face several inherent challenges, particularly in areas such as reasoning, consistency, and information’s contextual accuracy. These difficulties come from the probabilistic nature of LLMs, which can lead to hallucinations, lack of transparency, and challenges in handling structured data.

This is where Knowledge Graphs (KGs) come into play. By integrating LLMs with KGs, AI-generated knowledge can be significantly enhanced. Why? KGs provide a structured and interconnected representation of information, reflecting the relationships and entities in the real world. Unlike traditional databases, KGs can capture and reason about the complexities of human knowledge, ensuring that the outputs of LLMs come from a structured, verifiable knowledge base. This integration leads to more accurate, consistent, and contextually relevant outcomes.

Industries like healthcare, finance, and legal services can greatly benefit from knowledge graphs due to their need for precise and interpretable information.

Objective

The article delves into the innovative ways LLMs can be harnessed to build and leverage knowledge graphs.

After reading this article, you will understand:

1. How to build a knowledge graph using LLMs

2. How LLMs can then leverage knowledge graphs to enhance their capabilities

3. What are the limitations and key considerations when integrating LLMs with knowledge graphs?

No prior knowledge is required to understand the article. The experimentations described in the article were carried out focus only the libraries llama-index.

You can find the code here on GitHub.

Subscribe now

1. Navigating Knowledge Graphs & LLMs Landscape

This section provides a comprehensive overview of the process involved in building and leveraging knowledge graphs with LLMs. It introduces the key concepts, terminology, and various approaches before delving into the specific methods in the following sections.

1.1. How does the process of developing and leveraging knowledge graphs unfold?

Figure 1 illustrates the overall workflow, consisting of two key steps:

1. Knowledge graph extraction, along with the possible techniques detailed in Section 2. This step results in the creation of the knowledge graph.

2. Retrieval using one of the techniques outlined, which will be discussed an in-depth in Section 3. These techniques can be employed to answer queries.

1.2. Refresher on knowledge graph terminology

In network terminology, a triplet refers to a fundamental unit of a knowledge graph. It consists of three elements: a subject (an entity), a relation (the connection between entities), and an object (another entity).

Nodes represent the entities (e.g., people, companies) within the network, while edges or links represent the relationships or interactions between these nodes.

Together, nodes and edges form a graph structure, where triplets define the specific connections within this network.

1.3. About Data

The data for the experiments is drawn from the Wikipedia page on the Panama Papers scandal.

2. Building Knowledge Graphs: Extraction Methods Explored

The library llama-index proposes 4 main methods to extract a knowledge graph from a corpus of documents using LLMs:

1. Schema-based extraction

2. Free-form extraction

3. Dynamic extraction

4. Implicit extraction

Let’s delve into each of these methods.

2.1. Schema-based extraction

The main idea of this approach is to construct a graph based on a predefined network model. By specifying the desired entity types and their possible relationships, the LLM-based approach generates a network from the text, structured according to the specified model.

💡 Example: In the Panama Papers scandal, we can establish a predefined structure where entities include companies, banks, or scandals, and relationships are defined as ‘involved’ or ‘owns.’ The schema-based extracted graph would then be expected to reveal connections such as ‘Person A owns Company B’ or ‘Bank D is involved in Scandal E.’

How to implement schema-based extraction in practice?

Here’s a quick skeleton for implementing schema-based extraction in practice:

from typing import Literal
from llama_index.core.indices.property_graph import SchemaLLMPathExtractor
from llama_index.core import PropertyGraphIndex
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# Define the possible entity types for the knowledge graph
entities = Literal["PERSON", "COMPANY", "COUNTRY", "BANK", "SCANDAL"]

# Define the possible relations between the entities in the knowledge graph
relations = Literal["OWNS", "LOCATED_IN", "INVOLVED_IN"]

# Define the schema that outlines which entities can have which relations
schema = {
    "PERSON": ["OWNS", "LOCATED_IN", "INVOLVED_IN"],
    "COMPANY": ["OWNS", "LOCATED_IN", "INVOLVED_IN"],
    "COUNTRY": ["LOCATED_IN"],
    "BANK": ["LOCATED_IN", "INVOLVED_IN"],
    "SCANDAL": ["INVOLVED_IN"],
}

# Create an instance of SchemaLLMPathExtractor to extract paths based on the defined schema
kg_extractor = SchemaLLMPathExtractor(
  llm=OpenAI(model=LLM_MODEL, temperature=TEMPERATURE), # Use OpenAI's language model with the specified parameters
  possible_entities=entities,  # Define the types of entities to extract
  possible_relations=relations,  # Define the types of relations to extract
  kg_validation_schema=schema,  # Use the predefined schema for validation
  strict=True,  # Enforce strict validation; only entities and relations defined in the schema will be allowed
)

# Create a PropertyGraphIndex from the provided documents, using the specified embedding model
index = PropertyGraphIndex.from_documents(
    documents,  # The input documents to be processed and indexed
    embed_model=OpenAIEmbedding(model_name=EMBEDDING_MODEL),  # Use OpenAI's embedding model for document representation
    show_progress=True,  # Display progress during indexing
    kg_extractors=[kg_extractor],  # Use the previously defined SchemaLLMPathExtractor for extracting knowledge graph paths
)

# Define the storage path for the keyword extractor
path_output_storage_kg_extractor = f"{path_output_storage}/{kw_extractor_name}/"

# Create the storage directory if it doesn't already exist
if not os.path.exists(path_output_storage_kg_extractor):
    os.makedirs(path_output_storage_kg_extractor)

# Persist the index's storage context to the specified directory
index.storage_context.persist(persist_dir=path_output_storage_kg_extractor)

# Save the knowledge graph as a NetworkX graph to an HTML file
index.property_graph_store.save_networkx_graph(name=f"{path_output}/kwnoledge_graphs_{kw_extractor_name}.html")

⚠️ Note: In the kg_extractors argument of PropertyGraphIndex, multiple knowledge extractor instances can be added, enabling the combination of various extraction methods within a single network.

Figure 2 displays a part of the schema-based extracted graph generated by the code. The nodes represent either a person, company, country, bank, or scandal, while the relationships are limited to ownership and involvement, as predefined in the code.

What is the method behind schema-based extraction?

✓ Text analysis: The extractor first analyzes the text content within each chunk using an LLM. This model interprets the text to identify potential entities (such as people, organizations, or locations) and relationships (such as ownership, location, or association) based on the provided schema.

✓ Path extraction:

• Using the schema as a guide, the tool extracts specific relationships between identified entities, forming triplets.

• A triplet is a set of three components: a subject entity, a relationship, and an object entity (e.g., “Person X owns Company Y”).

• The max_triplets_per_chunk parameter limits the number of triplets processed at a time, ensuring efficiency when handling large datasets.

✓ Validation: The extracted triplets are then validated against the schema to ensure they fit the predefined types of entities and relationships. If the schema is set to strict mode, only triplets that fully comply with the schema are retained.

✓ Graph construction: The validated triplets are then used to construct or expand a knowledge graph. Each triplet adds nodes (entities) and edges (relationships) to the graph, gradually building a network that reflects the information extracted from the text.

✓ Metadata integration: The tool can also incorporate metadata, which may include additional context or properties associated with entities and relationships (e.g., timestamps, source documents). This helps in enriching the graph with more detailed and relevant information.

2.2. Free-form extraction

Free-form extraction is a flexible method of extracting relationships, or “triples,” from text data. Contrary to the scheme-based extraction methods, it doesn't not rely on a predefined schema or structure but leverages the capabilities of a language model (LLM) to identify meaningful connections between entities on its own.

💡 Example: In the context of the Panama Papers scandal, free-form extraction can analyze leaked documents to uncover hidden relationships between individuals, companies, and offshore entities. For instance, the model might extract triples such as “Company B is registered in Offshore Jurisdiction C”, or “Law Firm D manages Company B”. These connections, inferred from the context rather than explicitly stated, help build a knowledge graph that maps the complex network of entities involved in the scandal.

How to implement free-form extraction in practice?

from llama_index.core.indices.property_graph import SimpleLLMPathExtractor
# Create an instance of SimpleLLMPathExtractor
kg_extractor = SimpleLLMPathExtractor(
    llm=OpenAI(model=LLM_MODEL, temperature=TEMPERATURE)
)

Figure 3 presents a portion of the free-form extracted graph. Here, the relationships between nodes are more flexible, allowing for greater exploration. However, this flexibility involves a trade-off between the desire to explore and the need for a concise, essential-focused representation.

Source: “Former Emir of Qatar Hamad bin Khalifa Al Thani owned Afrodille S.A., which had a bank account in Luxembourg and shares in two South African companies. Al Thani also held a majority of the shares in Rienne S.A. and Yalis S.A., holding a term deposit with the Bank of China in Luxembourg. ” [3]

How does free-form extraction achieve its results?

✓ Input text analysis: The process begins with analyzing the input text from chunks. The content is passed through the language model, which uses a specific prompt to guide the extraction process. It generates potential triples. The extraction is controlled by the max_paths_per_chunk parameter, which limits the number of triples generated in each chunk. This ensures the process remains efficient and manageable, especially when handling large volumes of text.

✓ Parsing: The output from the language model is then parsed by a function that extracts these triples from the raw text response.

✓ Building the graph: The extracted triples are then used to create or update a knowledge graph. These nodes and edges are enriched with metadata, which can include additional information like the source of the data or timestamps.

2.3. Dynamic extraction

Dynamic extraction is similar to free-form extraction in its flexibility, but it introduces an initial structure that can evolve as the model processes more data.

While free-form extraction allows the model to interpret and extract relationships without any predefined structure, dynamic extraction starts with a basic ontology (ie. a set of initial entity types and relationships) that guides the extraction process.

This method not only extracts relationships but also encourages the discovery of new entity types and connections, expanding the ontology as needed.

💡 Example: In the Panama Papers scandal, dynamic extraction might start with an initial ontology that includes entities like “Person,” “Company,” and “Offshore Account,” and relationships such as “owns” or “controls.” As the LLM processes leaked documents, it could discover new relationships like “facilitated by” (connecting a person to a law firm) or identify new entity types such as “shell company.” These discoveries would expand the original ontology, allowing the knowledge graph to capture a broader and more nuanced picture of the entities and their connections in the scandal.

How to implement dynamic extraction in practice?

from llama_index.core.indices.property_graph import DynamicLLMPathExtractor
from llama_index.core import PropertyGraphIndex
from llama_index.llms.openai import OpenAI

# Define the possible entity types for the knowledge graph
entities = ["PERSON", "COMPANY", "COUNTRY", "BANK", "SCANDAL"]

# Define the possible relations between the entities in the knowledge graph
relations = ["OWNS", "LOCATED_IN", "INVOLVED_IN"]

# Create an instance of SimpleLLMPathExtractor
kg_extractor = DynamicLLMPathExtractor(
    llm=OpenAI(model=LLM_MODEL, temperature=TEMPERATURE),
    allowed_entity_types=entities,
    allowed_relation_types=relations,
    )# Define the possible entity types for the knowledge graph
entities = ["PERSON", "COMPANY", "COUNTRY", "BANK", "SCANDAL"]

Figure 4 presents a segment of the resulting dynamic extracted graph. It demonstrates that the approach successfully identified relevant new entities and relationships, building upon the initial ones provided. This technique strikes an ideal balance between exploration and structure.

How does the dynamic extraction method operate?

✓ Starting with an initial ontology: Unlike free-form extraction, which operates without a predefined schema, dynamic extraction begins with an initial ontology. This ontology includes basic entity types (e.g., “Person,” “Company”) and relationships (e.g., “owns,” “located in”). It serves as a starting point that the method can expand as it processes new data.

✓ Guided LLM prompting: Similar to free-form extraction, dynamic extraction uses prompts to instruct the language model (LLM) on what to extract. However, the prompts in dynamic extraction not only guide the extraction of known entities and relationships but also encourage the LLM to identify and propose new ones.

✓ Parsing and expanding the graph: The model’s output is parsed into triples (subject-relation-object) as in free-form extraction. However, in dynamic extraction, the process includes identifying and integrating new entity types or relationships.

✓ Ontology Evolution: As the model processes more text, it continuously refines and expands the ontology.

2.4. Implicit extraction

The key idea of this method is to use inherent connections within the data to identify relationships between entities, even when these relationships are not explicitly mentioned. It infers connections based on the structure or context of the data, making it useful when relationships can be logically deduced rather than directly stated.

💡 Example: Implicit extraction is most suited when analyzing a structured document with inherent relationships, such as a legal contract or a research paper with a clear hierarchy of sections, subsections, and references. In this case, the relationships between sections (e.g., “Chapter 1 is followed by Chapter 2”) or between references (e.g., “Section A cites Reference B”) are not explicitly labeled as relationships but can be inferred from the document’s structure. In that case, implicit extraction would automatically infer relationships like “Clause 1.1 is a subsection of Section 1” or “Clause 1.1 refers to Appendix A”.

How to implement implicit extraction in practice?

Here’s a quick skeleton for implementing implicit extraction in practice:

from llama_index.core.indices.property_graph import ImplicitPathExtractor

# Create an instance of SimpleLLMPathExtractor
kg_extractor = ImplicitPathExtractor()

How is implicit extraction carried out?

✓ Node relationships: Each node in a graph can have various types of relationships with other nodes, such as being the parent of another node, being sequentially before or after another node, or referencing another node as a source. These relationships are typically stored within the node’s structure.

✓ Edge extraction: The extractor scans through the list of nodes and examines these inherent relationships. For each node, it identifies connections to other nodes and creates corresponding edges (relations) between them. These edges are then labeled appropriately (e.g., “PARENT,” “CHILD,” “NEXT,” “PREVIOUS”) based on the type of relationship.

✓ Metadata integration: During this process, the extractor can also integrate any relevant metadata associated with the nodes, such as timestamps or source information, into the relationships it creates. This enriches the graph with additional context, making the implicit connections more informative.

3. Empowering LLMs through Knowledge Graph Synergy

Once the graph is built, the challenge becomes how to effectively leverage it for LLM queries. The Llama-Index library currently offers four approaches, two of which are cypher-based.

1. Vector context retriever

2. LLM synonym retriever

3. Text-to-cypher

4. Cypher template retriever

Cypher is a powerful query language used for interacting with graph databases, enabling precise retrieval and manipulation of graph data through pattern matching and other advanced querying techniques.”

3.1. Vector context retriever

This approach retrieves nodes based on their semantic similarity to a given query, allowing for more precise and contextually appropriate search results.

💡 Example: Imagine searching for information related to “offshore accounts” in a property graph built from the Panama Papers data. The vector context retriever would convert the query into a vector and find nodes in the graph that are semantically similar, such as nodes representing specific companies, individuals, or transactions linked to offshore accounts. It would then retrieve these nodes, along with any directly related entities (e.g., owners or connected banks), and return them sorted by relevance.

How can the Vector Context Retriever be applied in practice?

# Import necessary modules from the llama_index library
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# Import the VectorContextRetriever class from the property_graph module
from llama_index.core.indices.property_graph import VectorContextRetriever
# Create a sub-retriever using VectorContextRetriever
# This will use the property graph store and vector store from the loaded index
# The embed_model parameter specifies the model to be used for embedding queries (e.g., OpenAI's embedding model)
sub_retriever = VectorContextRetriever(
  index.property_graph_store, 
  vector_store=index.vector_store,
  embed_model=OpenAIEmbedding(model_name=EMBEDDING_MODEL),
)
# Create a retriever from the index using the previously defined sub-retriever
retriever = index.as_retriever(sub_retrievers=[sub_retriever])
# Initialize the query engine using the retriever
# The query engine will use the retriever(s) to process and return responses to queries
query_engine = index.as_query_engine(
    sub_retrievers=[retriever]
)
# Perform a query on the query engine to retrieve information about ICIJ's involvement in the Panama Papers scandal
print(
    query_engine.query(
    "How was the International Consortium of Investigative Journalists (ICIJ) involved in the Panama Papers scandal?"
    ).response
)

Output:
The International Consortium of Investigative Journalists (ICIJ) helped organize the research and document review once Süddeutsche Zeitung realized the scale of the work required to validate the authenticity of the leaked data. The ICIJ enlisted reporters and resources from various media outlets to investigate individuals and organizations associated with Mossack Fonseca. Additionally, the ICIJ released the leaked documents from the Panama Papers scandal on its website after verifying the source and content.

How does the Vector Context Retriever retriever work?

✓ Embedding the query: When a query is submitted, the vector context retriever first converts the query into a vector representation using an embedding model. This vector captures the semantic meaning of the query, enabling the retriever to find similar content within the graph.

✓ Vector store query: The retriever identifies the nodes whose vectors are most similar to the query vector, using a similarity measure (like cosine similarity). The number of top similar nodes to retrieve is controlled by the similarity_top_k parameter.

✓ Node retrieval and path expansion: Once the most similar nodes are identified, the retriever can also explore relationships connected to these nodes based on the specified path_depth.

✓ Similarity scoring: Each retrieved node is assigned a similarity score that reflects how closely it matches the query.

✓ Final output: The retriever sorts the retrieved nodes by their similarity scores, optionally filters them based on the similarity_score threshold, and then returns the most relevant nodes and their associated relationships.

3.2. LLM synonym retriever

The LLM synonym retriever enhances the retrieval of information from a knowledge graph by expanding a user’s query with synonyms generated by an LLM.

This expansion allows the retriever to capture and return more relevant nodes from the graph, improving the accuracy of the search results.

💡 Example: An investigator searching the Panama Papers graph for “offshore company” could use the LLM synonym retriever to expand the query with terms like “shell company” or “tax haven entity.” This broader search would uncover additional relevant nodes, such as companies involved in tax evasion, that might be missed with the original query alone, leading to a more thorough investigation.

How to use the LLM synonym retriever in practice?

# Import necessary modules from the llama_index library
from llama_index.core.indices.property_graph import LLMSynonymRetriever

sub_retriever = LLMSynonymRetriever(
    index.property_graph_store,  # The property graph store from the index is used as the data source
    llm=OpenAI(model=LLM_MODEL, temperature=TEMPERATURE),  # Initialize the LLM with specified model and temperature
    include_text=True,  # Include the source chunk text in the retrieved paths
    max_keywords=100,  # Maximum number of keywords to be considered for retrieval
    path_depth=5,  # Limit the depth of the search paths to 5 levels
)

What is the process behind the LLM synonym retriever?

✓ Query expansion with synonyms: When a query is submitted, the LLM synonym retriever first uses a language model to generate synonyms for the keywords in the query. This is done through a predefined prompt, which instructs the LLM to suggest alternative words or phrases similar to the original query terms.

✓ Synonym integration: The synonyms generated by the LLM are then integrated into the query. This allows the LLM to carry out a broader search within the graph.

✓ Graph retrieval:

• The expanded query is used to retrieve nodes from the graph. The retriever searches for nodes and relationships that match any of the terms in the expanded query.

• The depth of the retrieval, or how far the search goes in terms of relationships (e.g., direct connections or further down the relationship chain), can be adjusted using the path_depth parameter.

✓ Parsing and matching: The output from the LLM is parsed into a format that can be used to identify relevant nodes in the graph.

✓ Node scoring and retrieval: The matched nodes are scored based on their relevance to the original query and the relationships they hold within the graph. The retriever then returns a list of nodes with their associated scores, representing the most relevant entities and connections for the query.

3.3. Cypher-based retrievers

Llama-index offers two Cypher-based retrievers to leverage the capabilities of graph databases.

Here’s a brief overview without delving into the details:

• The TextToCypherRetriever operates by leveraging a graph store schema, an input query, and a prompt template for text-to-cypher conversion to generate and execute a cypher query.

• The CypherTemplateRetriever offers a more controlled approach compared to the TextToCypherRetriever. Instead of allowing the LLM to generate any cypher statement freely, it uses a predefined cypher template, where the LLM fills in the specific details.

Key Takeaways

✓ Knowledge graphs not only improve the accuracy and contextual relevance of AI-generated knowledge but also serve as a powerful tool for industries that demand precise and interpretable information, such as healthcare, finance, and legal services.

✓ When selecting a method to build it with LLMs, it’s essential to consider the trade-off between exploration and structure. Dynamic extraction offers an excellent balance for those who need both flexibility and structure in their knowledge graph.

✓ For leveraging a knowledge graph, various techniques are available, with vector-based methods providing efficient retrieval.

✓ However, it’s important to acknowledge that the current version of Llama-Index has limitations in network analysis and handling complex graph structures. These capabilities would certainly improve as the synergies between knowledge graphs and LLMs are a hot topic!

References

[1] Llama index documentation

[2] Introducing the Property Graph Index: A Powerful New Way to Build Knowledge Graphs with LLMs

[3] Panama Papers

Predictive analytics has long been a cornerstone of decision-making, but what if we told you there’s an alternative beyond forecasting? What if you could strategically influence the outcomes instead?

Uplift modeling holds this promise. It adds an interesting dynamic layer to traditional predictions by identifying individuals whose behavior can be influenced positively if they receive special treatments.

The application use cases are endless. In medicine, it would help identify patients for whom a medical treatment could improve their health. In retail, such a model allows for better targeting of customers for whom a promotion or personalized offering would be effective in retention.

Objective

This article is the first part of a series that explores the transformative potential of uplift modeling, shedding light on how it can reshape strategies in marketing, healthcare, and beyond. It focuses on uplift models based on decision trees and uses, as a case study, the prediction of customer conversion with the application of promotional offers

After reading this article, you will understand:

1. What exactly is uplift modeling?

2. In what ways can decision trees be tailored for uplift modeling?

3. How to assess the performance of uplift models?

No prior knowledge is required to understand the article.
The experimentations described in the article were carried out using the libraries scikit-uplift, causalml and plotly. You can find the code here on GitHub.

Subscribe now

1. Demystifying Uplift Modelling

1.1. Why uplift models?

The best way to understand the benefit of using uplift models is through an example. Imagine a scenario where a telecommunications company aims to reduce customer churn.

A “traditional” ML-based approach would consist of using a model trained on historical data to predict the likelihood of current customers to churn. This would help identify customers at risk and take proactive measures, such as offering tailored incentives or exclusive promotions.

• What if the customer is leaving regardless (for example, relocating to another country)?

• And what if the promotional offer sent by mail or over the phone reinforces the customer’s urge to change operator?

In such instances, even if the model accurately predicts their departure, the resulting actions may lead to wasted resources and expenses in communication. In the latter case, it could even contribute to an actual churn of customers who would have remained without action from the company.

This is where uplift modeling steps in, offering a more nuanced understanding of customer behavior. Its goal is not just to detect customers at risk but customers that will respond positively to a special treatment such as a promotion. Thus, businesses can accurately target the appropriate customers, streamlining resource allocation for optimal efficiency.

Uplift modeling refers to the set of techniques used to model the incremental impact of an action or treatment on a customer outcome. Source [1]

1.2. What to expect from uplift models?

Uplift models categorize individuals into different classes based on their responses to treatment. The main classes typically include:

1. Persuadables: Individuals who are positively influenced by the treatment. Without the intervention, they are less likely to respond as desired. Persuadables represent the group that benefits most from targeted interventions.

2. Sure Things: Individuals who respond positively to the treatment regardless of whether they receive it or not.

3. Lost Causes: Individuals who are unlikely to respond positively to the treatment, even if they receive it. Identifying this group helps avoid unnecessary resource allocation, as interventions are unlikely to yield the desired outcome.

4. Sleeping Dogs / Do-Not-Disturb: Individuals who respond negatively to the treatment, meaning the intervention might have a detrimental effect on their behavior. Recognizing this group is crucial to prevent inadvertently causing negative outcomes.

This classification is summarized in Figure 1. It plays an essential role in uplift modeling.

1.3. What data do uplift models need for training?

Data schema

Before diving into the theory, let’s look at the data needed to train uplift models:

• Feature columns: attributes or characteristics of the individuals or observations in the dataset that are used as input for the model.

• Treatment indicator: binary column that indicates whether a particular individual or observation has received a treatment or intervention.

• Target column: outcome variable that the model is designed to predict.

The experimentations carried out for the purpose of this article are based on Criteo dataset [8]. It was built by gathering the results of several incrementality tests. These tests follow a specific randomized trial approach in which a randomly selected portion of the population is deliberately excluded from receiving targeted advertising. The target variable is visit.

A rough approach to get an idea of the impact of treatment is to calculate the average treatment effect (ATE), i.e. the difference in the occurrence of visits between instances with and without treatment. In the Criteo dataset, the ATE is equal to 1.034%.

Hypothesis

Behind these data lie several assumptions that must hold for the model to be trained properly. The main ones include:

1. Random Assignment Hypothesis: The treatment assignment is random or quasi-random, ensuring that any differences in outcomes are attributable to the treatment and not pre-existing biases.

2. Unconfoundedness: The treatment assignment is not influenced by any other factors that may affect the potential outcome.

3. Stable Unit Treatment Value Assumption (SUTVA): The treatment assignment of one individual does not affect the outcome of another individual.

4. Positivity assumption: All individuals have a non-zero probability of receiving both the treatment and control conditions. Thus, there are no “unreachable” groups in the population.

5. No Hidden Confounding Variables Hypothesis: All relevant variables influencing the outcome are measured and included in the model. There are no hidden confounding variables that are unaccounted for.

6. Consistency: The treatment assignment is consistent across different populations and time periods. This means that the effect of the treatment on an individual unit is the same across different populations and periods.

2. Tailoring Tree-Based Models for Uplift Modelling

2.1. How is uplift defined?

To assess the incremental benefits of applying a treatment (such as offering exclusive promotions in our previous example), we need to compute the difference in outcomes between receiving the treatment and not receiving it.

This difference is called the causal effect [1] and can formulated as:

Where:

• W is equal to 1 if a person receives the treatment (treatment group), and 0 if he or she does not (control group)

• Y is the outcome

This definition raises a fundamental challenge: it is impossible to observe the outcomes of both scenarios simultaneously for the same individual.

Nevertheless, it is still possible to estimate the expected causal effect of the treatment for a subset of the population, called the Conditional Average Treatment Effect (CATE). It is expressed as:

Where X is the feature vector representing individual characteristics.

This estimation relies on the unconfoundedness assumption mentioned earlier, which translated to the assumption that the treatment assignment W is independent of the outcome Y given X. This condition holds when the treatment assignment is random conditional on X.

2.2. How are tree-based models adapted to uplift modeling?

Concepts related to uplift trees

Let’s start by uplift decision trees. They are very similar to conventional decision trees but with an important adjustment.

• Traditional decision trees are designed to maximize the accuracy of the predicted class probabilities. Their splitting criteria and pruning methods (Gini impurity or information gain, etc.) are adapted toward this objective.

• On the other hand, uplift decision trees are intended to model the change in class probabilities caused by treatment, such as mailing an offer or treating a patient. Thus, the splitting criteria and pruning methods are designed to maximize the differences between treatment and control probabilities in the left and right subtrees.

A perfect uplift decision tree would take the following form:

In that case, what splitting criteria should be used to maximize the differences between the treatment and control probabilities in the left and right subtrees?

P. Rzepakowski and S. Jaroszewicz introduced in [2] the following gain formulation for a given splitting choice (called test) A:

Where:

• D(.) is divergence measure

• PT(Y) and PC(Y) are the probability distribution of the outcome in respectively the treated group and control group

For the divergence metrics, they mentioned Kullback, Euclidean, and Chi-Squared, defined as:

To illustrate the computation of divergence gain, let’s consider an example:

Note that this model can model the causal effect of not just one but multiple treatments.

Training and inference process for tree-based uplift models

Let’s summarize the steps to train an uplift decision tree to identify the stages at which the introduced concepts come into play.

1. Building the uplift tree:

• Node splitting criteria: At each node, the algorithm calculates a measure of distribution divergence between the treatment and control groups. The algorithm chooses the feature and split point that maximize the distribution divergence, aiming to separate the treatment and control distributions.

• Recursive splitting: The tree-building process is recursive. At each node, the algorithm selects the best feature and split point based on the divergence gain and continues splitting until a stopping condition is met.

2. Stopping conditions: The algorithm typically stops when any of the specified conditions are satisfied:

• Maximum depth: Limit the depth of the tree to prevent overfitting.

• Minimum samples per leaf: Set a threshold for the minimum number of samples required to create a leaf node.

• Minimum improvement in divergence: Stop splitting a node if the gain in class distribution divergence uplift improvement is below a certain threshold.

3. Pruning: After building the tree, prune branches that do not contribute significantly to the model performance. This helps prevent overfitting and improves generalization to new data.

4. Uplift prediction:

• Leaf node assignment: For each observation, traverse the tree to determine the leaf node it falls into.

• Uplift score assignment: Calculate the predicted outcomes for the treatment and control groups separately at the leaf node. Assign the uplift score as the difference between the predicted outcomes for the treatment and control groups to capture the impact of the treatment on individual instances.

Similar to traditional random forests, uplift random forests leverage ensemble learning to improve performance. The approach involves training multiple decision trees and aggregating their predictions.

2.3. Application

To whet your appetite, here is an example of code that trains an uplift model and makes predictions. It relies on the library developed by Uber, CausalML [11]. The evaluation of uplift models is discussed in the rest of the article.

The complete code is available here on GitHub.

from causalml.inference.tree import UpliftRandomForestClassifier
# model initialization
# more information about the parameters in the library doc
# https://causalml.readthedocs.io/en/latest/causalml.html#causalml.inference.tree.UpliftRandomForestClassifier

uplift_model = UpliftRandomForestClassifier(
    n_estimators = 10,
    max_depth = 5,
    min_samples_leaf = 10,
    min_samples_treatment = 3,
    n_reg = 10,
    normalization = True,
    evaluationFunction = "ED",
    control_name='control',
    honesty = False
)

# model training
uplift_model.fit(
    train_df[cols].values,
    treatment=train_df['treatment_group'].values,
    y=train_df[col_target].values
)

# model inference
uplift_model.predict(
  test_df[cols].values, 
  full_output=True)[uplift_model.classes_]

# finding the best treatment
pred_df["best_treatment_index"] = np.argmax(pred_df.values, axis=1)
pred_df["best_treatment"] = pred_df["best_treatment_index"].map(
    {i:c for i, c in enumerate(uplift_model.classes_)}
)

Potential issues

One should be cautious about potential problems that can lead to an ill-trained model:

• Treatment imbalance issue: This happens when the algorithm concentrates most if not all, treatment instances in one subtree. It violates the unconfoundedness assumption which is crucial to uplift modeling. This situation makes further splits challenging and may lead to an ill-trained model.

• Predisposition for selecting splits with numerous child nodes: Such splits may yield higher gains in training data but tend to overfit when extrapolated to the test data.

To address these issues, normalization and regularization factors have been introduced in [2].

3. Quantifying Uplift Models Performance

3.1. What are the common evaluation metrics?

Evaluating uplift models is significantly different from traditional machine learning. Why? Since it’s impossible to observe the effects of being treated and untreated on an individual simultaneously, there is no ground truth.

Uplift by percentile

A common alternative is to calculate the uplift by decile. This involves the following steps:

1. Sort the individuals based on their predicted uplift scores in descending order

2. Divide the sorted dataset into percentiles containing a specific percentage of the population

3. Compute the uplift or response rate for each percentile by comparing the average in the treatment group with that in the control group within the percentile.

from sklift.metrics import uplift_by_percentile

uplift_by_perc_df =  uplift_by_percentile(
    test_df[col_target], 
    pred_df["uplift"],
    test_df["treatment_group"],
    strategy="overall", 
    total=False, 
    std=True, 
    bins=10
)

uplift_by_perc_df.head()

While this metric is straightforward, it does not allow comparison between models. What if one model performs better on the first decile but worse on the second decile compared to another model? Determining the best model is difficult in such cases. A simple idea is to calculate the cumulative uplift, but the approach still depends on the choice of decile size.

Uplift curve & AUUC

Uplift curve is a widespread approach to evaluating uplift models. It does not rely on predefined percentiles leaving threshold setting up to the evaluator.

It consists of the following steps:

1. Sort the individuals based on their predicted uplift scores in descending order

2. Compute the cumulative response rates for both the treatment group and the control group as we move down the sorted list. Thus, the cumulative response rate at each point is the fraction of responders in the group up to that point.

Here’s the typical shape of an uplift curve: the x-axis denotes the percentage of individuals targeted (e.g., top 1%, top 5%), while the y-axis illustrates the disparity in cumulative response rates between the treatment and control groups.

• Increasing uplift curve: A steeper incline suggests that the model is effectively identifying individuals who respond positively to the treatment. This often occurs in the early stages of the curve (left side). The greater the uplift at the beginning, the more impactful the model is for a smaller percentage of the population.

• Sustained growth: If the uplift curve continues to rise on the right side, it indicates that is it still effective to target individuals even if the percentage of this population is large.

• Flattening uplift curve: As the curve progresses, it may reach a point where the uplift plateaus or starts to flatten. This suggests that the model’s effectiveness in identifying responders diminishes as a larger portion of the population is considered. The uplift becomes less pronounced.

• Decreasing uplift curve: Conversely, if the uplift curve is decreasing, it signals diminishing returns. This happens as the model targets more individuals.

• Baseline: It is obtained using uniform random assignment. It is increasing when the average treatment effect is positive which is often the case.

The overall performance of the model can be assessed by calculating the Area Under the Uplift Curve (AUUC). A larger AUUC indicates a better-performing model across various percentages of the population.

from sklift.metrics import uplift_auc_score

# Area Under Uplift Curve
# link to the library doc:
# https://www.uplift-modeling.com/en/latest/api/metrics/uplift_auc_score.html

uplift_auc = uplift_auc_score(
    y_true=test_df[col_target], 
    uplift=pred_df["uplift"],
    treatment=test_df["treatment_group"],
)

Another option is to calculate the uplift at the top K or a specific percentage. This enables the selection of the best model, considering a constrained budget or resource allocation.

from sklift.metrics import uplift_at_k

# Uplift@30%
# link to the library doc:
# https://www.uplift-modeling.com/en/latest/api/metrics/uplift_at_k.html

uplift_at_k_score = uplift_at_k(
    y_true=test_df[col_target], 
    uplift=pred_df["uplift"],
    treatment=test_df["treatment_group"],
    strategy='overall', k=0.3
)

3.2. Application

The following section describes the experimentation carried out on Criteo dataset.

Methodology

• The dataset was divided into three subsets, maintaining an equal proportion of positive outcomes across each: one designated for training the model (constituting 70% of the total), another for testing (15%) to fine-tune hyperparameters, and a third for validation (15%) to facilitate model comparison.

• Model hyper-parameters were finetuing using first the AUUC as the metrics then upflit à 30%

The performance of the model was compared to:

• A baseline involving random assignment.

• A solo model [10] using a random forest as the estimator. This involves fitting a regular RF model with ‘treatment’ as an additional feature. To estimate uplift for an observation, the model is used twice: once with the treatment flag set to 1 and the other with 0. Uplift is then obtained by subtracting the results.

3.3. Results

This small experiment shows that the uplift model performs better than a solo model on the validation set: the AUUC is greater for the uplift model whose parameters have been fine-tuned to maximize upflit (in general rather than at 30%) on the test set.

Just as with the classic random forest, it is possible to visualize the most important variables for the uplift model.

Key Takeaways

✔️ Uplift modeling is a powerful analytical approach that aims to predict the incremental impact of a treatment or intervention on an individual’s behavior.

✔️ It is possible to adapt various machine learning techniques to uplift modeling, with random forests being a prime example. By adjusting the algorithm’s objective, one can capitalize on the inherent benefits of random forests, including good performance, robustness, interpretability, and easy implementation.

✔️ However, it is crucial to pay close attention to uplift modeling-specific assumptions. Rigorous checks are needed to ensure the model is well-trained.

✔️ Moreover, uplift modeling requires the interpretation of distinctive metrics, such as the uplift curve, which differ from traditional machine learning metrics. This is because there is no inherent ground truth for uplift by design.

✔️ Overall, uplift modeling goes beyond standard predictions. It provides insights into individuals who are likely to change their behavior if they receive a specific treatment, making it a powerful tool in reshaping strategies in marketing, healthcare, and beyond.”

References

[1] P. Gutierrez, J.Y. Gérardy, Causal Inference and Uplift Modeling, A review of the literature, 2016

[2] P. Rzepakowski, S. Jaroszewicz, Decision trees for uplift modeling with single and multiple treatments, 2011

[3] George Fei, Modeling Uplift Directly: Uplift Decision Tree with KL Divergence and Euclidean Distance as Splitting Criteria, 2019

[4] S. Athey, G.W. Imbens, Recursive Partitioning for Heterogeneous Causal Effects, 2015

[5] S. Athey and G.W. Imbens, Machine learning methods for estimating heterogeneous causal effects, 2015

[6] Ivan Klimuk, Uplift Modelling: Throw away your churn mode, Danske Technight

[7] F. Moraes, H.M. Proença, A. Kornilova, J. Albert, and D. Goldenberg, Uplift Modeling: From Causal Inference to Personalization, 32nd ACM International Conference on Information and Knowledge Management (CIKM ‘23), 2023

[8] Jing Pan, Why start using uplift models for more efficient marketing campaigns, Uber, 2020

[9] Kaggle, Uplift Modeling , Marketing Campaign Data, Criteo Dataset

[10] scikit-uplift documentation

[11] CausalML documentation

Recommending the right product to customers at the right time is a prevalent challenge across industries. For instance, bankers are constantly looking to suggest highly relevant services to their existing or potential customers. Retailers strive to recommend appealing products that meet customer tastes. Similarly, social networks aim to build captivating feeds to foster user adoption.

Despite being a widely explored use case, achieving satisfactory performance results remains arduous due to the unique nature of the problem. The main reasons include the presence of abundant categorical data, which often leads to scarcity issues, and the computational aspect of the use case, which poses scalability problems. It’s only recently that recommendation models have harnessed neural networks.

In this context, Meta has developed and made openly available a deep learning recommendation model (DRLM). The model is particularly remarkable for combining the principles of collaborative filtering and predictive analysis and being suitable for large-scale production.

Objective

The objective of this article is to guide you through a step-by-step implementation using the PyTorch library TorchRec, enabling you to effectively address your own recommendation use case.

After reading this article, you will understand:

1. How does the DLRM model work?

2. What sets DLRM models apart and makes them powerful and scalable?

3. How can you implement your own recommendation system from end to end?

The article requires general knowledge of the recommender system problem and familiarity with the pytorch library. The experimentations described in the article were carried out using the libraries TorchRec and PyTorch. You can find the code here on GitHub.

Subscribe now

1. Decoding the DLRM Model

Let’s first dive into the complexities of the DLRM model and explore its underlying principles and mechanisms.

1.1. An Overview of the Model Design

To provide a more tangible illustration, let’s consider the scenario of an online retailer looking to create a personalized feed for every customer visiting their website.

In order to achieve this, the retailer can train a model that predicts the probability of a customer purchasing a particular product. This model assigns a score to each product for each individual customer, based on various factors. The feed is built by ranking the scores.

In this case, the model can learn from historical data that encompasses a range of information for each customer and product. This includes numerical variables like customer age and product price, as well as categorical characteristics such as product type, color, and more.

Here’s where the DLRM model excels: it possesses the remarkable ability to leverage both numerical and categorical variables, even when dealing with a large number of unique categories. This enables the model to comprehensively analyze and understand complex relationships between the features. To understand why, let’s take a look at the architecture model in Figure 1.

Categorical features

DLRM learns an embedding table for each categorial feature and uses them to map these variables to dense representations. Hence, each categorical feature is represented as a vector of the same length.

Numerical features

DLRM processes numerical features through an MLP, called bottom MLP. The output of this MLP has the same dimension as the previous embedding vectors.

Pairwise interaction

DLRM calculates the dot product between all pairs of embedding vectors and the processed numerical features This allows the model to include second-order feature interaction.

Concatenation and final output

DLRM concatenates these dot products with the processed numerical features and uses the results to feed another MLP, called the top MLP. The final probability is obtained by passing the output of this MLP to a sigmoid function.

1.2. The Model Implementation

While the potential of the model appears promising in theory, its practical implementation presents a computational hurdle.

Usually, recommendation use cases involve handling vast volumes of data. Using DLRM models, in particular, introduces an very large number of parameters, higher than common deep learning models. Consequently, this amplifies the computational demands associated with their implementation.

1. The majority of parameters in DLRMs can be attributed to embeddings as they consist of multiple tables, each demanding a large memory. This makes DLRMs computationally demanding, both in terms of memory capacity and bandwidth.

2. Although the memory footprint of MLP parameters is smaller, they still require substantial computational resources.

To mitigate the memory bottleneck, DLRM relies a unique combination of model parallelism for the embeddings and data parallelism for the MLPs.

2. From Concept to Implementation: A Step-by-Step Guide to Building Your Custom Recommendation System

This section provides a detailed step-by-step guide on how to implement your own recommendation system from start to finish.

2.1. Data Transformation and Batch Construction

The first step involves converting the data into tensors and organizing them into batches for input into the model.

To illustrate this process, let’s consider this dataframe as an example.

For sparse features, we need to concatenate the values into a single vector and compute the lengths. This can be accomplished using the KeyedJaggedTensor.from_lengths_sync function, which takes both elements as input. Here’s an example of the Python script:

values = sample[cols_sparse].sum(axis=0).sum(axis=0)
values = torch.tensor(values).to(device)
# values = tensor([1, 0, 2, 0, 2, 2, 0, 2, 0, 1, 0, 1, 2, 0], device='cuda:0')

lengths = torch.tensor(
    pd.concat([sample[feat].apply(lambda x: len(x)) for feat in cols_sparse],
              axis=0).values,
    dtype=torch.int32
).to(self.device)
# lengths = tensor([1, 1, 1, 1, 1, 2, 3, 2, 2, 0], device='cuda:0', dtype=torch.int32)

sparse_features = KeyedJaggedTensor.from_lengths_sync(
  keys=cols_sparse,
  values=values,
  lengths=lengths
)

For dense features and labels, the process is more straightforward. Here’s an example of the Python script:

dense_features = torch.tensor(sample[cols_dense].values, dtype=torch.float32).to(device)
labels = torch.tensor(sample[col_label].values, dtype=torch.int32).to(device)

By using the outputs from the previous steps, it becomes possible to construct a batch. Here’s an example of the Python script:

batch = Batch(
  dense_features=dense_features,
  sparse_features=sparse_features,
  labels=labels,
).to(device)

For a more comprehensive implementation, you can refer to the batch.py file in the corresponding GitHub repository.

2.2. Model Initialization and Optimization Setup

The next step involves initializing the model, as demonstrated in the following Python code:

# Initialize the model and set up optimization

# Define the dimensionality of the embeddings used in the model
embedding_dim = 10

# Calculate the number of embeddings per feature
num_embeddings_per_feature = {c: len(v) for c, v in map_sparse.items()}

# Define the layer sizes for the dense architecture
dense_arch_layer_sizes = [512, 256, embedding_dim]

# Define the layer sizes for the overall architecture
over_arch_layer_sizes = [512, 512, 256, 1]

# Specify whether to use Adagrad optimizer or SGD optimizer
adagrad = False

# Set the epsilon value for Adagrad optimizer
eps = 1e-8

# Set the learning rate for optimization
learning_rate = 0.01

# Create a list of EmbeddingBagConfig objects for each sparse feature
eb_configs = [
    EmbeddingBagConfig(
        name=f"t_{feature_name}",
        embedding_dim=embedding_dim,
        num_embeddings=num_embeddings_per_feature[feature_name + '_enc'],
        feature_names=[feature_name + '_enc'],
    )
    for feature_idx, feature_name in enumerate(cols_sparse)
]

# Initialize the DLRM model with the embedding bag collection and architecture specifications
dlrm_model = DLRM(
    embedding_bag_collection=EmbeddingBagCollection(
        tables=eb_configs, device=device
    ),
    dense_in_features=len(cols_dense),
    dense_arch_layer_sizes=dense_arch_layer_sizes,
    over_arch_layer_sizes=over_arch_layer_sizes,
    dense_device=device,
)

# Create a DLRMTrain instance for handling training operations
train_model = DLRMTrain(dlrm_model).to(device)

# Choose the appropriate optimizer class for the embedding parameters
embedding_optimizer = torch.optim.Adagrad if adagrad else torch.optim.SGD

# Set the optimizer keyword arguments
optimizer_kwargs = {"lr": learning_rate}
if adagrad:
    optimizer_kwargs["eps"] = eps

# Apply the optimizer to the sparse architecture parameters
apply_optimizer_in_backward(
    optimizer_class=embedding_optimizer,
    params=train_model.model.sparse_arch.parameters(),
    optimizer_kwargs=optimizer_kwargs,
)

# Initialize the dense optimizer with the appropriate parameters
dense_optimizer = KeyedOptimizerWrapper(
    dict(in_backward_optimizer_filter(train_model.named_parameters())),
    optimizer_with_params(adagrad, learning_rate, eps),
)

# Create a CombinedOptimizer instance to handle optimization
optimizer = CombinedOptimizer([dense_optimizer])

The model can then be trained and evaluated using the following code:

loss, (loss2, logits, labels) = train_model(batch)

For a more comprehensive implementation, you can refer to the model.py file in the corresponding GitHub repository.

Key Takeaways

✔ The DLRM model presents a compelling approach to effectively combine numerical and categorical features using embeddings, enabling the model to capture intricate patterns and relationships.

✔ Although its architecture requires considerable computational resources, its implementation incorporates a unique combination of model parallelism and data parallelism, making the model scalable to production.

✔ However, due to limited data availability, the model’s performance has not been extensively tested across diverse real-world datasets. This raises uncertainty about its effectiveness in practical scenarios.

✔ Additionally, the model necessitates tuning a considerable number of parameters, further complicating the process.

✔ Considering this, simpler models like LGBM may offer comparable performance with easier implementation, tuning, and long-term maintenance, without the same computational overhead.

References

[1] M Naumov & al, Deep Learning Recommendation Model for Personalization and Recommendation Systems, May 2019

[2] Github repository of the Facebook team’s initial implementation of the DLRM model available as open source

[3] DLRM: An advanced, open source deep learning recommendation model, Meta AI Blog, July 2019

[4] Pytorch library for modern production recommendation systems, torchec

[5] Vinh Nguyen, Tomasz Grel and Mengdi Huang, Optimizing the Deep Learning Recommendation Model on NVIDIA GPUs, June 2020

Survival models are great for predicting the time for an event to occur. These models can be used in a wide variety of use cases including predictive maintenance (forecasting when a machine is likely to break down), marketing analytics (anticipating customer churn), patient monitoring (predicting a patient is likely to be re-hospitalized), and much more.

By combining machine learning with survival models, the resulting models can benefit from the high predictive power of the former while retaining the framework and typical outputs of the latter (such as the survival probability or hazard curve over time). For more information, check out the first article of this series here.

However, in practice, ML-based survival models still require extensive feature engineering and thus prior business knowledge and intuition to lead to satisfying results. So, why not use deep learning models instead to bridge the gap?

Objective

This article focuses on how deep learning can be combined with the survival analysis framework to solve use cases such as predicting the likelihood of a patient being (re)hospitalized.

After reading this article, you will understand:

1. How can deep learning be leveraged for survival analysis?

2. What are the common deep learning models in survival analysis and how do they work?

3. How can these models be applied concretely to hospitalization forecasting?

This article is the second part of the series around survival analysis. If you are not familiar with survival analysis, it is best to start by reading the first one here. The experimentations described in the article were carried out using the libraries scikit-survival, pycox, and plotly. You can find the code here on GitHub.

Subscribe now

1. Survival Analysis & Deep Learning: How Can They Be Combined?

1.1. Problem statement

Let’s start by describing the problem at hand.

We are interested in predicting the likelihood that a given patient will be rehospitalized given the available information about his health status. More specifically, we would like to estimate this probability at different time points after the last visit. Such an estimate is essential to monitor patient health and mitigate their risk of relapse.

This is a typical survival analysis problem. The data consists of 3 elements:

Patient’s baseline data including:

• Demographics: age, gender, locality (rural or urban)

• Patient history: smoking, alcohol, diabetes mellitus, hypertension, etc.

• Laboratory results: hemoglobin, total lymphocyte count, platelets, glucose, urea, creatinine, etc.

• More information about the source dataset here.

A time t and an event indicator δ∈{0;1}:

• If the event occurs during the observation duration, t is equal to the time between the moment the data were collected and the moment the event (i.e., rehospitalization) is observed, In that case, δ = 1.

• If not, t is equal to the time between the moment the data were collected and the last contact with the patient (e.g. end of study). In that case, δ = 0.

⚠️ With this description, why use survival analysis methods when the problem is so similar to a regression task? The initial paper gives a pretty good explanation of the main reason:

“If one chooses to use standard regression methods, the right-censored data becomes a type of missing data. It is usually removed or imputed, which may introduce bias into the model. Therefore, modeling right-censored data requires special attention, hence the use of a survival model.” Source [2]

1.2. DeepSurv

Approach

Let’s move on to the theoretical part with a little refresher on the hazard function.

“The hazard function is the probability an individual will not survive an extra infinitesimal amount of time δ, given they have already survived up to time t. Thus, a greater hazard signifies a greater risk of death.”

Source [2]

Similar to the Cox proportional hazards (CPH) model, DeepSurv is based on the assumption that the hazard function is the product of the 2 functions:

• the baseline hazard function: λ_0(t)

• the risk score, r(x)=exp(h(x)). It models how the hazard function varies from the baseline for a given individual given the observed covariates.

More on CPH models in the first article of this series.

The function h(x) is commonly referred to as the log-risk function. And this is precisely the function that the Deep Surv model aims at modeling.

In fact, CPH models assume that h(x) is a linear function: h(x) = β . x. Fitting the model consists thus in computing the weights β to optimize the objective function. However, the linear proportional hazards assumption does not hold in many applications. This justifies the need for a more complex non-linear model that is ideally capable of handling large volumes of data.

Architecture

In this context, how can the DeepSurv model provide a better alternative? Let’s start by describing it. According to the original paper, it’s a “deep feed-forward neural network which predicts the effects of a patient’s covariates on their hazard rate parameterized by the weights of the network θ.” [2]

How does it work?

‣ The input to the network is the baseline data x.

‣ The network propagates the inputs through a number of hidden layers with weights θ. The hidden layers consist of fully-connected nonlinear activation functions followed by dropout.

‣ The final layer is a single node that performs a linear combination of the hidden features. The output of the network is taken as the predicted log-risk function.

Source [2]

As a result of this architecture, the model is very flexible. Hyperparametric search techniques are typically used to determine the number of hidden layers, the number of nodes in each layer, the dropout probability and other settings.

What about the objective function to optimize?

• CPH models are trained to optimize the Cox partial likelihood. It consists of calculating for each patient i at time Ti the probability that the event has happened, considering all the individuals still at risk at time Ti, and then multiplying all these probabilities together. You can find the exact mathematical formula here [2].

• Similarly, the objective function of DeepSurv is the log-negative mean of the same partial likelihood with an additional part that serves to regularize the network weights. [2]

Code sample

Here is a small code snippet to get an idea of how this type of model is implemented using the pycox library. The complete code can be found in the notebook examples of the library here [6].

# Step 1: Neural net
# simple MLP with two hidden layers, ReLU activations, batch norm and dropout

in_features = x_train.shape[1]
num_nodes = [32, 32]
out_features = 1
batch_norm = True
dropout = 0.1
output_bias = False

net = tt.practical.MLPVanilla(in_features, num_nodes, out_features, batch_norm,
                              dropout, output_bias=output_bias)

model = CoxPH(net, tt.optim.Adam)


# Step 2: Model training

batch_size = 256
epochs = 512
callbacks = [tt.callbacks.EarlyStopping()]
verbose = True

model.optimizer.set_lr(0.01)

log = model.fit(x_train, y_train, batch_size, epochs, callbacks, verbose,
                val_data=val, val_batch_size=batch_size)


# Step 3: Prediction

_ = model.compute_baseline_hazards()
surv = model.predict_surv_df(x_test)


# Step 4: Evaluation

ev = EvalSurv(surv, durations_test, events_test, censor_surv='km')
ev.concordance_td()

1.3. DeepHit

Approach

Instead of making strong assumptions about the distribution of survival times, what if we could train a deep neural network that would learn them directly?

This is the case with the DeepHit model. In particular, it brings two significant improvements over previous approaches:

• It does not rely on any assumptions about the underlying stochastic process. Thus, the network learns to model the evolution over time of the relationship between the covariates and the risk.

• It can handle competing risks (e.g., simultaneously modeling the risks of being rehospitalized and dying) through a multi-task learning architecture.

Architecture

As described here [3], DeepHits follows the common architecture of multi-task learning models. It consists of two main parts:

1. A shared subnetwork, where the model learns from the data a general representation useful for all the tasks.

2. Task-specific subnetworks, where the model learns more task-specific representations.

However, the architecture of the DeepHit model differs from typical multi-task learning models in two aspects:

• It includes a residual connection between the inital covariates and the input of the task-specific sub-networks.

• It uses only one softmax output layer. Thanks to this, the model does not learn the marginal distribution of competing events but the joint distribution.

The figures below show the case where the model is trained simultaneously on two tasks.

The output of the DeepHit model is a vector y for every subject. It gives the probability that the subject will experience the event k ∈ [1, 2] for every timestamp t within the observation time.

2. Use Case Application: How Do These Models Perform in Practice?

2.1. Methodology

Data

The data set was divided into three parts: a training set (60% of the data), a validation set (20%), and a test set (20%). The training and validation sets were used to optimize the neural networks during training and the test set for final evaluation.

Benchmark

The performance of the deep learning models was compared to a benchmark of models including CoxPH and ML-based survival models (Gradient Boosting and SVM). More information on these models is available in the first article of the series.

Metrics

Two metrics were used to evaluate the models:

• Concordance index (C-index): it measures the capability of the model to provide a reliable ranking of survival times based on individual risk scores. It is computed as the proportion of concordant pairs in a dataset.

• Brier score: It’s a time-dependent extension of the mean squared error to right censored data. In other words, it represents the average squared distance between the observed survival status and the predicted survival probability.

2.2. Results

In terms of C-index, the performance of the deep learning models is considerably better than that of the ML-based survival analysis models. Moreover, there is almost no difference between the performance of Deep Surval and Deep Hit models.

In terms of Brier score, the Deep Surv model stands out from the others.

• When examining the curve of the Brier score as a function of time, the curve of the Deep Surv model is lower than the others, which reflects a better accuracy.

• This observation is confirmed when considering the integration of the score over the same time interval.

Note that the Brier wasn’t computed for the SVM as this score is only applicable for models that are able to estimate a survival function.

Finally, deep learning models can be used for survival analysis as well as statistical models. Here, for instance, we can see the survival curve of randomly selected patients. Such outputs can bring many benefits, in particular allowing a more effective follow-up of the patients that are the most at risk.

Key Takeaways

✔️ Survival models are very useful for predicting the time it takes for an event to occur.

✔️ They can help address many use cases by providing a learning framework and techniques as well as useful outputs such as the probability of survival or the hazard curve over time.

✔️ They are even indispensable in this type of uses cases to exploit all the data including the censored observations (when the event does not occur during the observation period for example).

✔️ ML-based survival models tend to perform better than statistical models (more information here). However, they require high-quality feature engineering based on solid business intuition to achieve satisfactory results.

✔️ This is where Deep Learning can bridge the gap. Deep learning-based survival models like DeepSurv or DeepHit have the potential to perform better with less effort!

✔️ Nevertheless, these models are not without drawbacks. They require a large database for training and require fine-tuning multiple hyperparameters.

References

[1] Bollepalli, S.C.; Sahani, A.K.; Aslam, N.; Mohan, B.; Kulkarni, K.; Goyal, A.; Singh, B.; Singh, G.; Mittal, A.; Tandon, R.; Chhabra, S.T.; Wander, G.S.; Armoundas, A.A. An Optimized Machine Learning Model Accurately Predicts In-Hospital Outcomes at Admission to a Cardiac Unit. Diagnostics 2022, 12, 241.

[2] Katzman, J., Shaham, U., Bates, J., Cloninger, A., Jiang, T., & Kluger, Y. (2016). DeepSurv: Personalized Treatment Recommender System Using A Cox Proportional Hazards Deep Neural Network, ArXiv

[3] Laura Löschmann, Daria Smorodina, Deep Learning for Survival Analysis, Seminar information systems (WS19/20), February 6, 2020

[4] Lee, Changhee et al. DeepHit: A Deep Learning Approach to Survival Analysis With Competing Risks. AAAI Conference on Artificial Intelligence (2018).

[5] Wikipedia, Proportional hazards model

[6] Pycox library

Predicting the probability of an event occurring is good, predicting the time remaining before an event occurs is even better!

Take the example of customer churn. What if, instead of predicting the probability of a customer leaving the company in the next months, you could predict this probability at several time points over the next months? The benefits of such an approach are immediate. It would allow you to anticipate and prioritize your marketing actions more effectively in time and, ultimately, reduce the churn rate.

This falls exactly in the field of survival analysis, also called time-to-event analysis. It refers to a learning framework and a set of techniques that can be used to estimate the time it takes for an event of interest to occur based on observations.

The name of survival analysis comes from the typical use case where it was first applied: predicting time to death for clinical research. However, one should not be misled by its name: it is not limited to the medical field, but can be applied to use cases in multiple industries. And with the recent advances in data science, survival analysis has reemerged leaving the world of classical statistics to include more advanced methods in machine learning.

Objective

This article focuses on how machine learning can be combined with survival analysis framework to solve use cases such as predicting churn.

After reading this article, you will understand:

1. What is survival analysis all about?

2. What are the main survival models and how do they work?

3. How can these models be applied concretely to churn prediction?

This article is the first part of the series around survival analysis. No prior knowledge is required to understand the article. The experimentations described in the article were carried out using the libraries scikit-survival and plotly. You can find the code here on GitHub.

Subscribe now

1. What is Survival Analysis About?

1.1. Problem statement

At first glance, survival analysis may appear to be just another regression problem since the goal is to predict when an event is likely to occur (a continuous variable). However, the problem has a twist: a part of the training data may be partially observed — it is censored.

To illustrate this, let’s take the example of a company that offers services on a subscription basis. The company wants to predict for each customer that reaches out to the support the probability of unsubscribing over time. During the period in which the data was collected:

• Customer A did not unsubscribe before the end of the study.

• Customers B and D canceled their subscriptions after a few months.

• Customer C has decided to restrict access to its data from the platform.

In this case, the records of customers A and C are censored.

More formally, each observation contains a set of covariates X = (x_1, …, x_n), the time t when the event occurred, or the censoring time c>0. Let’s introduce an event indicator δ∈{0;1}. The observable time y of a right-censored sample is then defined as:

In our case of churn prediction, the data consists of the customer’s support contact history. Each observation contains information about:

• the interaction: the date, the reason for the call (sign-up/support), and the channel (email/phone).

• the customer: age and gender.

• subscription: product, price, billing cycle (monthly/annual), sign-up.

The data was enriched with additional features including the number of times the customer has contacted the company in the past, the duration since the customer subscribed, and cyclical date-related features.

1.2. Use cases

Survival analysis can be used for a wide range of applications where the goal is to predict the time between two events. Here are few other examples:

• Predictive maintenance: Forecast when the machine is likely to break down after it has been turned on. Data can be censored if the machines had to stop due to external factors (fire alarms, etc.)

• Patient monitoring: Predict when a patient is likely to be re-hospitalized since their first diagnosis or hospitalization. Data may be censored if the patient has left the geographic scope of the study.

• Marketing analytics: Predict when a prospect might convert since the first call. Data may be censored if the person has died during the observation period.

• Economics: Forecast the time it takes for a laid-off person to find a job. Data may be censored if a person drops out of the study.

2. What Are Common Approaches to Tackle Survival Use Cases?

2.1. Kaplan — Meier estimator

One of the most simple approaches to this problem is to use the Kaplan-Meier estimator. This is a non-parametric method that focuses on approximating the survival function. Let’s start by discussing the fundamental underlying theories before examining the results of the application case.

Survival function

The survival function S(t) represents the probability that a subject survives beyond time t or, in a similar way, the probability that the duration is at least equal to t.

where T is the random lifespan taken from the study population.

S starts with 1 at time t=0 since no subject has yet experienced the event at the beginning. It decreases and tends to 0 as everyone is likely to experience the event of interest at some point.

Kaplan-Meier estimator

To approximate the survival function, the Kaplan-Meier model breaks the estimation into small steps. For each interval, the probability is computed as follows:

where n_i is the number of individuals who are at risk at time point t_i and d_i the number of individuals that experienced the event at time t_i.

It is a very simple method that does not take into account covariates.

• It can be used as a naive baseline.

• It can also serve as a data exploration method. In this case, it provides an overview of the survival function for the whole population or helps in comparing certain groups between them.

For example, in our case study, we might compare the estimate depending on the subscription billing cycle. The graph below confirms the intuition that customers with a monthly subscription are more volatile. They tend to churn more often and faster during the first years after their subscription.

2.2. Cox proportional hazards model

The most widely used estimator is certainly the Cox proportional hazards (Cox PH) model. It is easy to implement, takes covariates into account and provides interpretable results. It is a semi-parametric method that aims at modeling the hazard function.

Hazard function

The hazard function h(t) represents the probability of the death event occurring at time t, given that the subject did not experience the death event until time t.