Gradient Health Monitoring

Relevant source files

The following files were used as context for generating this wiki page:

• neuraldbg.py
• demo_vanishing_gradients.py
• infrastructure_planning/CODEBASE_GUIDE.md
• README.md
• plans/reasoning_engine_analysis.md

Introduction

Gradient Health Monitoring is a core subsystem of the NeuralDBG engine designed to observe, classify, and analyze the stability of gradients during deep learning training. Unlike traditional logging tools that focus on raw metrics, this system extracts semantic events—high-level transitions in learning dynamics—to identify exactly when and where training instabilities like vanishing or exploding gradients originate.

The system operates non-invasively by attaching PyTorch hooks to model layers, allowing it to monitor the backward pass without requiring modifications to the user's training loop. By tracking gradient norms across temporal steps, it builds a causal trace that helps researchers understand the root causes of training failures.

Sources: README.md, infrastructure_planning/CODEBASE_GUIDE.md

Architecture and Core Components

The monitoring architecture follows a "Hub & Spokes" design where the core logic resides in neuraldbg.py (the Hub) and interacts with the PyTorch framework via hooks (a Spoke).

Semantic Classification

The system categorizes gradient stability using a specialized enumeration. Gradients are classified based on their norm values relative to predefined heuristics.

Health Status	Description	Threshold Logic (Heuristic)
HEALTHY	Gradients are in an optimal range for learning.	$10^{-6} \le \text{norm} \le 1.0$
VANISHING	Gradients are too small, preventing weight updates.	$\text{norm} < 10^{-6}$
EXPLODING	Gradients are excessively large, causing instability.	$\text{norm} > 10^3$
SATURATED	High norm values indicating potential saturation.	$\text{norm} > 1.0$

Sources: neuraldbg.py:25-30, plans/reasoning_engine_analysis.md:Section 7.1

Causal Inference Engine

The NeuralDbg class serves as the main entry point, managing the lifecycle of gradient observation. It utilizes a context manager to automate hook installation and removal.

The diagram above illustrates the operational flow from hook installation to the generation of semantic events based on health status changes.
Sources: neuraldbg.py, infrastructure_planning/CODEBASE_GUIDE.md

Logic and Data Flow

Event Detection Mechanism

Gradient health monitoring relies on detecting transitions. A SemanticEvent is triggered when the gradient health status of a specific layer changes between training steps. This allows the system to perform Causal Compression, focusing only on meaningful shifts rather than every numerical update.

1. _classify_gradient_health(norm): Categorizes the raw tensor norm into a GradientHealth enum.
2. _detect_gradient_transition(layer_id, current_health): Compares the current health against the previously recorded state for that specific layer.
3. Semantic Event Creation: If a transition (e.g., HEALTHY → VANISHING) is detected, an event is logged with metadata including the step number, layer ID, and confidence score.

Sources: neuraldbg.py:300-350, infrastructure_planning/CODEBASE_GUIDE.md

Temporal and Structural Coupling

The engine analyzes how failures propagate through the network. It identifies if a vanishing gradient in layer $N$ is coupled with a similar failure in layer $N-1$, helping to identify the "Root Cause" layer.

The sequence diagram details the interaction between the PyTorch backward pass and the NeuralDBG engine during a single training step.
Sources: neuraldbg.py, plans/reasoning_engine_analysis.md:Section 1

Implementation Example: Vanishing Gradients

The system is specifically tuned to detect vanishing gradients, a common failure in deep networks using Tanh or Sigmoid activations. In demo_vanishing_gradients.py, the engine monitors a deep Tanh network to demonstrate this capability.

# From demo_vanishing_gradients.py
with NeuralDbg(model) as dbg:
    for step, (x, y) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output, y)
        loss.backward() # Hooks trigger here
        optimizer.step()

# Post-mortem analysis
explanations = dbg.explain_failure()

The explain_failure() method performs abductive reasoning over the collected gradient events to rank potential causes, such as learning rate mismatches or activation saturation.
Sources: demo_vanishing_gradients.py:80-120, neuraldbg.py:400-450

Key Data Structures

Structure	Purpose
GradientHealth	Enum defining the state of gradients (HEALTHY, VANISHING, etc.).
SemanticEvent	Data class representing a detected transition in the training dynamics.
CausalHypothesis	Result of the reasoning engine, providing a ranked explanation of a failure.

Sources: neuraldbg.py:20-60, infrastructure_planning/CODEBASE_GUIDE.md

Summary

Gradient Health Monitoring in NeuralDBG transforms raw numerical training data into a structured causal narrative. By utilizing non-invasive hooks and a transition-based event model, the system identifies the precise step and layer where gradient instabilities begin. This allows developers to move beyond simple metric tracking to true post-mortem causal analysis of neural network training failures.