Introduction

Variational Autoencoders (VAEs) represent a groundbreaking fusion of deep learning and Bayesian inference, standing as one of the most influential architectures in generative modeling. Unlike traditional autoencoders that simply compress and reconstruct data, VAEs learn a probabilistic latent space, enabling them to generate new, realistic data samples by sampling from this learned distribution. This capability makes VAEs invaluable for tasks where data generation, anomaly detection, and representation learning are critical. In enterprise environments, the challenge extends beyond building a single VAE model. It requires orchestrating multiple intelligent agents that can collaborate, maintain context through persistent memory, and retrieve relevant domain knowledge using Graph RAG (Retrieval-Augmented Generation). In this article, we will build a complete proof-of-concept demonstrating how VAEs integrate into an enterprise-grade multi-agent system using LangGraph. Our use case focuses on financial transaction anomaly detection, where the VAE identifies unusual patterns, and a multi-agent workflow investigates, contextualizes, and reports these anomalies with explainable insights.

Understanding VAEs in Generative Modeling

A VAE consists of two main components:

  1. Encoder: Maps input data to a probability distribution in the latent space (typically a Gaussian distribution defined by mean μ and variance σ).

  2. Decoder: Reconstructs data from samples drawn from this latent distribution.

The key innovation is the reparameterization trick, which allows backpropagation through the stochastic sampling process. The loss function combines:

  • Reconstruction Loss: Ensures the output resembles the input (e.g., Mean Squared Error).

  • KL Divergence: Regularizes the latent space to follow a standard normal distribution, ensuring smoothness and generative capability.

This probabilistic approach enables VAEs to not only reconstruct data but also generate new samples by sampling from the latent space, making them powerful tools for generative modeling.

Real-Time Use Case: Financial Transaction Anomaly Detection & Investigation

Consider a large bank processing millions of transactions daily. Fraudulent activities often manifest as subtle deviations from normal behavior. A VAE can learn the distribution of "normal" transactions. When a new transaction is processed, the VAE attempts to reconstruct it. A high reconstruction error indicates an anomaly.

Our multi-agent system enhances this basic detection:

  1. VAE Agent: Processes transactions, calculates reconstruction error, and flags anomalies.

  2. Context Retriever Agent: Uses Graph RAG to fetch related historical cases, customer profiles, and known fraud patterns.

  3. Investigator Agent: Analyzes the anomaly score, retrieved context, and generates an investigative report with risk assessment.

  4. Memory Manager: Maintains state across interactions, learns from investigator feedback, and updates the knowledge graph.

Step-by-Step Implementation

Step 1: Setting Up the Environment

# requirements.txt
langgraph==0.2.0
langchain==0.1.0
torch==2.1.0
scikit-learn==1.3.0
fastapi==0.109.0
uvicorn==0.27.0
pydantic==2.5.0
chromadb==0.4.22
networkx==3.2.1
numpy==1.24.0

Step 2: Building the VAE Model

# models/vae.py
import torch
import torch.nn as nn
import torch.nn.functional as F

class VAE(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super(VAE, self).__init__()
        
        # Encoder
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        
        # Decoder
        self.fc3 = nn.Linear(latent_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, input_dim)
        
    def encode(self, x):
        h = F.relu(self.fc1(x))
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        h = F.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h))  # Assuming normalized input [0,1]
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon_x = self.decode(z)
        return recon_x, mu, logvar

def vae_loss(recon_x, x, mu, logvar):
    """Combines BCE reconstruction loss and KL divergence"""
    recon_loss = F.binary_cross_entropy(recon_x, x, reduction='sum')
    kl_divergence = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_divergence

Step 3: Implementing Graph RAG Service

# services/graph_rag.py
import chromadb
import networkx as nx
from typing import List, Dict

class GraphRAGService:
    def __init__(self):
        self.client = chromadb.PersistentClient(path="./chroma_db")
        self.collection = self.client.get_or_create_collection("fraud_cases")
        self.graph = nx.DiGraph()
    
    def add_case(self, case_id: str, description: str, metadata: Dict):
        """Add fraud case to vector store and knowledge graph"""
        self.collection.add(
            documents=[description],
            ids=[case_id],
            metadatas=[metadata]
        )
        self.graph.add_node(case_id, **metadata)
        if metadata.get('linked_case'):
            self.graph.add_edge(case_id, metadata['linked_case'])
    
    def retrieve_similar_cases(self, query: str, n_results: int = 3) -> List[Dict]:
        """Retrieve similar historical fraud cases"""
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results
        )
        return [{
            "id": results['ids'][0][i],
            "description": results['documents'][0][i],
            "metadata": results['metadatas'][0][i]
        } for i in range(len(results['ids'][0]))]
    
    def get_related_entities(self, case_id: str) -> List[str]:
        """Get linked entities from knowledge graph"""
        if self.graph.has_node(case_id):
            return list(self.graph.neighbors(case_id))
        return []

Step 4: Creating the LangGraph Multi-Agent Workflow

# agents/workflow.py
from langgraph.graph import StateGraph, END
from typing import TypedDict, List, Optional
import torch
from models.vae import VAE, vae_loss
from services.graph_rag import GraphRAGService

class AgentState(TypedDict):
    transaction_data: List[float]  # Normalized transaction features
    anomaly_score: Optional[float]
    reconstructed_data: Optional[List[float]]
    retrieved_context: List[dict]
    investigation_report: Optional[str]
    risk_level: Optional[str]  # Low, Medium, High
    conversation_history: List[dict]

class FraudDetectionAgent:
    def __init__(self):
        # Load pre-trained VAE (assuming input_dim=10, hidden=64, latent=2)
        self.vae = VAE(input_dim=10, hidden_dim=64, latent_dim=2)
        self.vae.load_state_dict(torch.load('models/vae_fraud.pth'))
        self.vae.eval()
        self.rag_service = GraphRAGService()
        self.threshold = 0.5  # Anomaly threshold
    
    def detect_anomaly(self, state: AgentState) -> AgentState:
        """VAE Agent Node: Calculate reconstruction error"""
        with torch.no_grad():
            x = torch.tensor([state['transaction_data']], dtype=torch.float32)
            recon_x, mu, logvar = self.vae(x)
            
            # Calculate per-sample reconstruction error
            error = torch.mean((x - recon_x) ** 2, dim=1).item()
            state['anomaly_score'] = error
            state['reconstructed_data'] = recon_x.squeeze().tolist()
            
            return state
    
    def retrieve_context(self, state: AgentState) -> AgentState:
        """Context Retriever Agent Node"""
        if state['anomaly_score'] > self.threshold:
            query = f"Transaction anomaly score {state['anomaly_score']:.4f}"
            state['retrieved_context'] = self.rag_service.retrieve_similar_cases(query)
        else:
            state['retrieved_context'] = []
        return state
    
    def investigate(self, state: AgentState) -> AgentState:
        """Investigator Agent Node: Generate report"""
        score = state['anomaly_score']
        if score > self.threshold:
            if score > 0.8:
                state['risk_level'] = "High"
                state['investigation_report'] = f"CRITICAL: High anomaly detected (Score: {score:.4f}). Review immediately. Context: {len(state['retrieved_context'])} similar cases found."
            else:
                state['risk_level'] = "Medium"
                state['investigation_report'] = f"WARNING: Moderate anomaly (Score: {score:.4f}). Monitor closely."
        else:
            state['risk_level'] = "Low"
            state['investigation_report'] = f"Normal transaction (Score: {score:.4f}). No action required."
        
        state['conversation_history'].append({
            "role": "system",
            "content": state['investigation_report']
        })
        return state

def build_workflow():
    """Build the LangGraph workflow"""
    agent = FraudDetectionAgent()
    
    workflow = StateGraph(AgentState)
    
    workflow.add_node("detect", agent.detect_anomaly)
    workflow.add_node("retrieve", agent.retrieve_context)
    workflow.add_node("investigate", agent.investigate)
    
    workflow.set_entry_point("detect")
    workflow.add_edge("detect", "retrieve")
    workflow.add_edge("retrieve", "investigate")
    workflow.add_edge("investigate", END)
    
    return workflow.compile()

Step 5: FastAPI Backend

# main.py
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from agents.workflow import build_workflow

app = FastAPI(title="VAE Fraud Detection System")
workflow = build_workflow()

class TransactionRequest(BaseModel):
    features: List[float]  # Normalized transaction features
    history: List[dict] = []

class DetectionResponse(BaseModel):
    anomaly_score: float
    risk_level: str
    report: str
    similar_cases: List[dict]

@app.post("/detect", response_model=DetectionResponse)
async def detect_fraud(request: TransactionRequest):
    initial_state = {
        "transaction_data": request.features,
        "anomaly_score": None,
        "reconstructed_data": None,
        "retrieved_context": [],
        "investigation_report": None,
        "risk_level": None,
        "conversation_history": request.history
    }
    
    result = await workflow.ainvoke(initial_state)
    
    return DetectionResponse(
        anomaly_score=result['anomaly_score'],
        risk_level=result['risk_level'],
        report=result['investigation_report'],
        similar_cases=result['retrieved_context']
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

Step 6: Simple Frontend Interface

<!-- index.html -->
<!DOCTYPE html>
<html>
<head>
    <title>VAE Fraud Detector</title>
    <style>
        body { font-family: Arial; max-width: 800px; margin: 50px auto; padding: 20px; }
        .result { padding: 15px; margin-top: 20px; border-radius: 5px; }
        .high { background: #ffcccc; border: 1px solid red; }
        .medium { background: #fff3cd; border: 1px solid orange; }
        .low { background: #d4edda; border: 1px solid green; }
        button { background: #007bff; color: white; padding: 10px 20px; border: none; cursor: pointer; }
    </style>
</head>
<body>
    <h1>VAE-Based Fraud Detection System</h1>
    <p>Enter 10 normalized transaction features (comma-separated):</p>
    <input type="text" id="features" placeholder="0.1, 0.5, 0.3, ..." style="width: 100%; padding: 10px;">
    <br><br>
    <button onclick="detectFraud()">Analyze Transaction</button>
    <div id="result" class="result" style="display:none;"></div>

    <script>
        async function detectFraud() {
            const input = document.getElementById('features').value;
            const features = input.split(',').map(Number);
            
            if (features.length !== 10 || features.some(isNaN)) {
                alert("Please enter exactly 10 numeric values.");
                return;
            }
            
            const response = await fetch('/detect', {
                method: 'POST',
                headers: {'Content-Type': 'application/json'},
                body: JSON.stringify({features: features})
            });
            const data = await response.json();
            
            const resultDiv = document.getElementById('result');
            resultDiv.className = `result ${data.risk_level.toLowerCase()}`;
            resultDiv.innerHTML = `
                <h3>Risk Level: ${data.risk_level}</h3>
                <p>Anomaly Score: ${data.anomaly_score.toFixed(4)}</p>
                <p><strong>Report:</strong> ${data.report}</p>
                ${data.similar_cases.length > 0 ? `<h4>Similar Cases:</h4><ul>${data.similar_cases.map(c => `<li>${c.description.substring(0, 100)}...</li>`).join('')}</ul>` : ''}
            `;
            resultDiv.style.display = 'block';
        }
    </script>
</body>
</html>

Conclusion

This proof-of-concept demonstrates the power of combining Variational Autoencoders with enterprise-grade multi-agent orchestration using LangGraph. The VAE provides a robust, probabilistic foundation for anomaly detection by learning the underlying distribution of normal data. When integrated with Graph RAG for contextual retrieval and a stateful multi-agent workflow for investigation, the system becomes a powerful tool for real-time financial fraud detection.

  • VAEs excel at learning latent representations and detecting anomalies through reconstruction error.

  • LangGraph enables seamless orchestration of specialized agents (Detector, Retriever, Investigator).

  • Graph RAG enhances decision-making by providing relevant historical context and entity relationships.

  • State Management ensures continuity and traceability throughout the investigation process.

This architecture is scalable, explainable, and ready for production deployment, showcasing how generative AI principles can be applied to critical enterprise security challenges.