Causal Compressor
Relevant source files
The following files were used as context for generating this wiki page:
The Causal Compressor is a fundamental architectural component of the NeuralDBG engine designed to handle the complexity of training dynamics by identifying patterns and propagation in training failures. Its primary function is to transform thousands of low-level Semantic Events into a concise set of root causes through a process of structural and temporal reduction.
By operating at module boundaries, the Causal Compressor ensures compatibility with optimizations like torch.compile while maintaining a high-level view of how instabilities, such as vanishing gradients or activation regime shifts, propagate through network layers.
Sources: infrastructure_planning/CODEBASE_GUIDE.md, README.md
Core Responsibilities
The Causal Compressor serves as the intermediary between raw event detection and the final abductive reasoning stage. It performs three primary operations:
- First Occurrence Identification: Pinpointing the exact layer and training step where a specific failure originated.
- Propagation Pattern Detection: Tracking how a failure in one layer (e.g., layer 3) triggers downstream or upstream failures in other layers.
- Coupled Failure Detection: Identifying temporal and structural links between different types of events, such as how an activation saturation event might be coupled with a subsequent gradient vanishing event.
Sources: infrastructure_planning/CODEBASE_GUIDE.md, plans/reasoning_engine_analysis.md, README.md
Data Flow Architecture
The following diagram illustrates how the Causal Compressor integrates into the broader NeuralDBG reasoning pipeline:
The Causal Compressor reduces the volume of data by focusing on the transition points and the causal links between events rather than the raw metric values.
Sources: plans/reasoning_engine_analysis.md:12-25, infrastructure_planning/CODEBASE_GUIDE.md
Compression Mechanisms
Temporal and Structural Coupling
The compressor identifies "coupled failures" by analyzing the proximity of events in both time (training steps) and space (network architecture). Currently, the system uses simple temporal proximity—typically within a 5-step window—to suggest a relationship between events.
| Method | Logic | Goal |
|---|---|---|
get_root_causes() | Identifies the first layer and step of failure. | Isolate the origin point. |
detect_coupled_failures() | Identifies temporal and structural links. | Find correlated instabilities. |
trace_causal_chain() | Links events across different layers. | Map the propagation path. |
Sources: infrastructure_planning/CODEBASE_GUIDE.md, plans/reasoning_engine_analysis.md:300-315, README.md
Semantic Event Structure
The Causal Compressor operates on SemanticEvent objects. These objects represent meaningful transitions rather than raw tensor data, allowing the compressor to reason about "health states."
# Conceptual representation of event metadata used for compression
# Sources: [README.md], [infrastructure_planning/CODEBASE_GUIDE.md]
{
"type": "gradient_health", # e.g., stable -> vanishing
"layer_id": "linear1",
"step_range": (234, 240),
"confidence": 0.87,
"causal_metadata": {
"propagation": "upstream",
"coupled_events": ["activation_saturation_layer2"]
}
}
Implementation within the NeuralDbg Engine
The Causal Compressor's logic is embedded within the NeuralDbg core class. It utilizes the results of state classification (e.g., ActivationHealth and GradientHealth) to detect transitions.
The compressor monitors shifts between states like NORMAL, SATURATED, and DEAD for activations, and HEALTHY, VANISHING, or EXPLODING for gradients.
Sources: infrastructure_planning/CODEBASE_GUIDE.md, neuraldbg.py
Roadmap and Future Enhancements
As part of the project roadmap, the Causal Compressor is slated for several upgrades to move beyond simple heuristic-based coupling:
- Phase 1 (Foundation): Implementation of event pruning and garbage collection to handle long training runs by keeping only high-confidence or recent events.
- Phase 3 (Causal Reasoning): Introduction of a Structural Causal Model (SCM) to explicitly define known relationships (e.g., how high learning rates cause exploding gradients) and transition from correlation-based detection to temporal precedence analysis.
- Performance Optimization: Integration of circular buffers (e.g.,
collections.deque) to limit event storage and improve memory efficiency.
Sources: ROADMAP.md, plans/reasoning_engine_analysis.md:340-365
Conclusion
The Causal Compressor is essential for distilling the "semantic trace" of learning dynamics into actionable insights. By identifying the root cause and the propagation chain of training failures, it allows NeuralDBG to provide ranked causal hypotheses that go beyond simple metric tracking provided by traditional tools.
Sources: README.md, infrastructure_planning/CODEBASE_GUIDE.md