trustformers_debug/visualization/
gradient_animation.rs

1//! Animated gradient flow visualization.
2//!
3//! Records per-layer gradient statistics across training steps and produces
4//! exportable frame sequences (JSON, CSV) and ASCII heatmap animations.
5
6use anyhow::{Context, Result};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::io::Write as _;
10use std::path::Path;
11
12// ============================================================================
13// Data types
14// ============================================================================
15
16/// Statistics for a single layer captured at one training step.
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct LayerGradientInfo {
19    /// Layer / parameter name.
20    pub name: String,
21    /// Mean of absolute gradient values.
22    pub mean_abs_grad: f64,
23    /// Maximum absolute gradient value.
24    pub max_abs_grad: f64,
25    /// L2 norm of the gradient vector.
26    pub grad_norm: f64,
27    /// `true` when the mean absolute gradient is below the vanishing threshold.
28    pub is_vanishing: bool,
29    /// `true` when the maximum absolute gradient exceeds the exploding threshold.
30    pub is_exploding: bool,
31    /// Normalised flow intensity in `[0.0, 1.0]` used for visual colour encoding.
32    pub flow_intensity: f64,
33}
34
35/// One frame of the gradient animation — all layers at a single training step.
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct GradientFrame {
38    /// The global training step index.
39    pub step: u64,
40    /// Per-layer information for this frame.
41    pub layers: Vec<LayerGradientInfo>,
42}
43
44// ============================================================================
45// Health classification
46// ============================================================================
47
48/// Overall gradient health of the training run.
49#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
50pub enum GradientHealth {
51    Healthy,
52    MinorIssues,
53    ProblemsDetected,
54    Critical,
55}
56
57/// Summary report produced by `GradientFlowAnimator::summary_report()`.
58#[derive(Debug, Clone)]
59pub struct GradientFlowSummary {
60    /// Number of steps recorded.
61    pub total_steps: u64,
62    /// Layers that exhibited vanishing gradients in at least one frame.
63    pub layers_with_vanishing_grads: Vec<String>,
64    /// Layers that exhibited exploding gradients in at least one frame.
65    pub layers_with_exploding_grads: Vec<String>,
66    /// Overall health classification for the run.
67    pub overall_health: GradientHealth,
68    /// Actionable recommendations derived from the observed gradient behaviour.
69    pub recommendations: Vec<String>,
70}
71
72// ============================================================================
73// Animator
74// ============================================================================
75
76/// Collects per-step gradient data and provides export / analysis utilities.
77pub struct GradientFlowAnimator {
78    frames: Vec<GradientFrame>,
79    /// Maximum number of frames to retain (rolling window, oldest dropped first).
80    max_frames: usize,
81    /// Mean absolute gradient below this value is flagged as vanishing.
82    vanishing_threshold: f64,
83    /// Maximum absolute gradient above this value is flagged as exploding.
84    exploding_threshold: f64,
85}
86
87impl GradientFlowAnimator {
88    /// Create a new animator.
89    ///
90    /// * `max_frames` — rolling window size (0 means unlimited).
91    pub fn new(max_frames: usize) -> Self {
92        Self {
93            frames: Vec::new(),
94            max_frames,
95            vanishing_threshold: 1e-7,
96            exploding_threshold: 1e3,
97        }
98    }
99
100    /// Set a custom vanishing-gradient threshold (default: 1e-7).
101    pub fn with_vanishing_threshold(mut self, threshold: f64) -> Self {
102        self.vanishing_threshold = threshold;
103        self
104    }
105
106    /// Set a custom exploding-gradient threshold (default: 1e3).
107    pub fn with_exploding_threshold(mut self, threshold: f64) -> Self {
108        self.exploding_threshold = threshold;
109        self
110    }
111
112    /// Record gradient tensors for one training step.
113    ///
114    /// `gradients` maps a layer name to the flat gradient vector for that layer.
115    pub fn record_step(&mut self, step: u64, gradients: &HashMap<String, Vec<f64>>) {
116        // Compute global max norm across all layers to normalise `flow_intensity`.
117        let global_max_norm: f64 = gradients.values().map(|g| l2_norm(g)).fold(0.0_f64, f64::max);
118
119        let mut layers: Vec<LayerGradientInfo> = gradients
120            .iter()
121            .map(|(name, grad)| {
122                let mean_abs = mean_abs(grad);
123                let max_abs = max_abs(grad);
124                let norm = l2_norm(grad);
125                let is_vanishing = mean_abs < self.vanishing_threshold;
126                let is_exploding = max_abs > self.exploding_threshold;
127                let flow_intensity = if global_max_norm > 0.0 {
128                    (norm / global_max_norm).clamp(0.0, 1.0)
129                } else {
130                    0.0
131                };
132                LayerGradientInfo {
133                    name: name.clone(),
134                    mean_abs_grad: mean_abs,
135                    max_abs_grad: max_abs,
136                    grad_norm: norm,
137                    is_vanishing,
138                    is_exploding,
139                    flow_intensity,
140                }
141            })
142            .collect();
143
144        // Stable ordering for deterministic output.
145        layers.sort_by(|a, b| a.name.cmp(&b.name));
146
147        let frame = GradientFrame { step, layers };
148        self.frames.push(frame);
149
150        // Enforce rolling window.
151        if self.max_frames > 0 && self.frames.len() > self.max_frames {
152            self.frames.remove(0);
153        }
154    }
155
156    /// All retained frames (possibly a rolling subset).
157    pub fn frames(&self) -> &[GradientFrame] {
158        &self.frames
159    }
160
161    /// Export all frames to a JSON file.
162    pub fn export_json(&self, path: &Path) -> Result<()> {
163        let json = serde_json::to_string_pretty(&self.frames)
164            .context("failed to serialise gradient frames")?;
165        if let Some(parent) = path.parent() {
166            std::fs::create_dir_all(parent).with_context(|| {
167                format!("failed to create output directory: {}", parent.display())
168            })?;
169        }
170        std::fs::write(path, json).with_context(|| {
171            format!(
172                "failed to write gradient animation JSON: {}",
173                path.display()
174            )
175        })?;
176        Ok(())
177    }
178
179    /// Export a CSV timeline: `step,layer,mean_abs_grad,max_abs_grad,grad_norm`.
180    pub fn export_csv(&self, path: &Path) -> Result<()> {
181        if let Some(parent) = path.parent() {
182            std::fs::create_dir_all(parent).with_context(|| {
183                format!("failed to create output directory: {}", parent.display())
184            })?;
185        }
186
187        let mut file = std::fs::File::create(path)
188            .with_context(|| format!("failed to create CSV file: {}", path.display()))?;
189
190        writeln!(
191            file,
192            "step,layer,mean_abs_grad,max_abs_grad,grad_norm,is_vanishing,is_exploding"
193        )
194        .context("failed to write CSV header")?;
195
196        for frame in &self.frames {
197            for layer in &frame.layers {
198                writeln!(
199                    file,
200                    "{},{},{:.8e},{:.8e},{:.8e},{},{}",
201                    frame.step,
202                    layer.name,
203                    layer.mean_abs_grad,
204                    layer.max_abs_grad,
205                    layer.grad_norm,
206                    layer.is_vanishing as u8,
207                    layer.is_exploding as u8,
208                )
209                .context("failed to write CSV row")?;
210            }
211        }
212        Ok(())
213    }
214
215    /// Render an ASCII "heatmap" animation string suitable for terminal display.
216    ///
217    /// Each row represents one layer; each column one recorded step.
218    /// Intensity is encoded with the characters `' ', '░', '▒', '▓', '█'`.
219    pub fn to_ascii_animation(&self) -> String {
220        if self.frames.is_empty() {
221            return "(no gradient frames recorded)\n".to_string();
222        }
223
224        // Collect all unique layer names in stable order.
225        let layer_names: Vec<String> = {
226            let mut seen: HashMap<&str, ()> = HashMap::new();
227            let mut names: Vec<String> = Vec::new();
228            for frame in &self.frames {
229                for layer in &frame.layers {
230                    if seen.insert(layer.name.as_str(), ()).is_none() {
231                        names.push(layer.name.clone());
232                    }
233                }
234            }
235            names.sort();
236            names
237        };
238
239        let blocks = [' ', '░', '▒', '▓', '█'];
240        let max_name_len = layer_names.iter().map(|n| n.len()).max().unwrap_or(8);
241
242        let mut out = String::new();
243        out.push_str("Gradient Flow Animation (step → right, layer ↓)\n");
244        out.push_str(&format!("{:>width$}  ", "layer", width = max_name_len));
245        for (i, _) in self.frames.iter().enumerate() {
246            out.push_str(&format!("{}", i % 10));
247        }
248        out.push('\n');
249        out.push_str(&"─".repeat(max_name_len + 2 + self.frames.len()));
250        out.push('\n');
251
252        for layer_name in &layer_names {
253            out.push_str(&format!("{:>width$}  ", layer_name, width = max_name_len));
254            for frame in &self.frames {
255                let intensity = frame
256                    .layers
257                    .iter()
258                    .find(|l| l.name == *layer_name)
259                    .map(|l| l.flow_intensity)
260                    .unwrap_or(0.0);
261                let idx = ((intensity * (blocks.len() - 1) as f64).round() as usize)
262                    .min(blocks.len() - 1);
263                out.push(blocks[idx]);
264            }
265            out.push('\n');
266        }
267
268        out
269    }
270
271    /// Generate a summary report for the recorded gradient history.
272    pub fn summary_report(&self) -> GradientFlowSummary {
273        let total_steps = self.frames.last().map(|f| f.step + 1).unwrap_or(0);
274
275        let mut vanishing: HashMap<String, ()> = HashMap::new();
276        let mut exploding: HashMap<String, ()> = HashMap::new();
277
278        for frame in &self.frames {
279            for layer in &frame.layers {
280                if layer.is_vanishing {
281                    vanishing.insert(layer.name.clone(), ());
282                }
283                if layer.is_exploding {
284                    exploding.insert(layer.name.clone(), ());
285                }
286            }
287        }
288
289        let mut layers_with_vanishing_grads: Vec<String> = vanishing.into_keys().collect();
290        layers_with_vanishing_grads.sort();
291        let mut layers_with_exploding_grads: Vec<String> = exploding.into_keys().collect();
292        layers_with_exploding_grads.sort();
293
294        let overall_health = classify_health(
295            &layers_with_vanishing_grads,
296            &layers_with_exploding_grads,
297            &self.frames,
298        );
299
300        let recommendations = build_recommendations(
301            &overall_health,
302            &layers_with_vanishing_grads,
303            &layers_with_exploding_grads,
304        );
305
306        GradientFlowSummary {
307            total_steps,
308            layers_with_vanishing_grads,
309            layers_with_exploding_grads,
310            overall_health,
311            recommendations,
312        }
313    }
314}
315
316// ============================================================================
317// Internal helpers
318// ============================================================================
319
320fn mean_abs(values: &[f64]) -> f64 {
321    if values.is_empty() {
322        return 0.0;
323    }
324    values.iter().map(|v| v.abs()).sum::<f64>() / values.len() as f64
325}
326
327fn max_abs(values: &[f64]) -> f64 {
328    values.iter().map(|v| v.abs()).fold(0.0_f64, f64::max)
329}
330
331fn l2_norm(values: &[f64]) -> f64 {
332    values.iter().map(|v| v * v).sum::<f64>().sqrt()
333}
334
335fn classify_health(
336    vanishing: &[String],
337    exploding: &[String],
338    frames: &[GradientFrame],
339) -> GradientHealth {
340    // Count frames with any issue.
341    let issue_frames = frames
342        .iter()
343        .filter(|f| f.layers.iter().any(|l| l.is_vanishing || l.is_exploding))
344        .count();
345    let total = frames.len().max(1);
346    let issue_ratio = issue_frames as f64 / total as f64;
347
348    if !exploding.is_empty() && issue_ratio > 0.5 {
349        return GradientHealth::Critical;
350    }
351    if !exploding.is_empty() || issue_ratio > 0.3 {
352        return GradientHealth::ProblemsDetected;
353    }
354    if !vanishing.is_empty() || issue_ratio > 0.1 {
355        return GradientHealth::MinorIssues;
356    }
357    GradientHealth::Healthy
358}
359
360fn build_recommendations(
361    health: &GradientHealth,
362    vanishing: &[String],
363    exploding: &[String],
364) -> Vec<String> {
365    let mut recs = Vec::new();
366
367    if !vanishing.is_empty() {
368        recs.push(format!(
369            "Vanishing gradients detected in: {}. Consider residual connections, layer normalisation, or a larger learning rate.",
370            vanishing.join(", ")
371        ));
372        recs.push("Investigate weight initialisation — Xavier or Kaiming init can prevent early vanishing.".to_string());
373    }
374    if !exploding.is_empty() {
375        recs.push(format!(
376            "Exploding gradients detected in: {}. Apply gradient clipping (clip_grad_norm).",
377            exploding.join(", ")
378        ));
379        recs.push("Consider reducing the learning rate or switching to a gradient-friendly optimiser (e.g. AdamW with weight decay).".to_string());
380    }
381    match health {
382        GradientHealth::Critical => {
383            recs.push("CRITICAL: Training stability is severely compromised — halt training and diagnose before continuing.".to_string());
384        },
385        GradientHealth::ProblemsDetected => {
386            recs.push("Significant gradient issues detected. Review architecture depth and learning rate schedule.".to_string());
387        },
388        GradientHealth::MinorIssues => {
389            recs.push("Minor gradient issues detected. Monitor closely; intervention may not be required immediately.".to_string());
390        },
391        GradientHealth::Healthy => {
392            recs.push("Gradients appear healthy — no immediate action required.".to_string());
393        },
394    }
395
396    recs
397}
398
399// ============================================================================
400// Tests
401// ============================================================================
402
403#[cfg(test)]
404mod tests {
405    use super::*;
406    use std::env::temp_dir;
407
408    fn simple_grads(layer_names: &[&str], value: f64) -> HashMap<String, Vec<f64>> {
409        layer_names
410            .iter()
411            .map(|&n| (n.to_string(), vec![value, value, value, value]))
412            .collect()
413    }
414
415    #[test]
416    fn test_record_step_basic() {
417        let mut animator = GradientFlowAnimator::new(0);
418        let grads = simple_grads(&["attn", "ffn"], 0.01);
419        animator.record_step(0, &grads);
420        assert_eq!(animator.frames().len(), 1);
421        let frame = &animator.frames()[0];
422        assert_eq!(frame.layers.len(), 2);
423    }
424
425    #[test]
426    fn test_rolling_window() {
427        let mut animator = GradientFlowAnimator::new(3);
428        for i in 0..10u64 {
429            let grads = simple_grads(&["layer_a"], 0.1);
430            animator.record_step(i, &grads);
431        }
432        assert_eq!(
433            animator.frames().len(),
434            3,
435            "rolling window should cap at max_frames"
436        );
437    }
438
439    #[test]
440    fn test_vanishing_detection() {
441        let mut animator = GradientFlowAnimator::new(0).with_vanishing_threshold(1e-6);
442        let mut grads = HashMap::new();
443        grads.insert("shallow".to_string(), vec![1e-8, 1e-8]);
444        grads.insert("deep".to_string(), vec![0.01, 0.01]);
445        animator.record_step(0, &grads);
446        let frame = &animator.frames()[0];
447        let shallow = frame.layers.iter().find(|l| l.name == "shallow").unwrap();
448        assert!(shallow.is_vanishing);
449        let deep = frame.layers.iter().find(|l| l.name == "deep").unwrap();
450        assert!(!deep.is_vanishing);
451    }
452
453    #[test]
454    fn test_exploding_detection() {
455        let mut animator = GradientFlowAnimator::new(0).with_exploding_threshold(100.0);
456        let mut grads = HashMap::new();
457        grads.insert("bad_layer".to_string(), vec![500.0, 200.0]);
458        grads.insert("ok_layer".to_string(), vec![0.1, 0.1]);
459        animator.record_step(0, &grads);
460        let frame = &animator.frames()[0];
461        let bad = frame.layers.iter().find(|l| l.name == "bad_layer").unwrap();
462        assert!(bad.is_exploding);
463        let ok = frame.layers.iter().find(|l| l.name == "ok_layer").unwrap();
464        assert!(!ok.is_exploding);
465    }
466
467    #[test]
468    fn test_flow_intensity_normalised() {
469        let mut animator = GradientFlowAnimator::new(0);
470        let mut grads = HashMap::new();
471        grads.insert("large".to_string(), vec![1.0; 10]);
472        grads.insert("small".to_string(), vec![0.001; 10]);
473        animator.record_step(0, &grads);
474        let frame = &animator.frames()[0];
475        for layer in &frame.layers {
476            assert!(layer.flow_intensity >= 0.0 && layer.flow_intensity <= 1.0);
477        }
478    }
479
480    #[test]
481    fn test_export_json() -> Result<()> {
482        let mut animator = GradientFlowAnimator::new(0);
483        animator.record_step(0, &simple_grads(&["a", "b"], 0.1));
484        animator.record_step(1, &simple_grads(&["a", "b"], 0.05));
485
486        let path = temp_dir().join(format!("grad_anim_{}.json", uuid::Uuid::new_v4()));
487        animator.export_json(&path)?;
488        assert!(path.exists());
489        let content = std::fs::read_to_string(&path)?;
490        let frames: Vec<GradientFrame> = serde_json::from_str(&content)?;
491        assert_eq!(frames.len(), 2);
492        Ok(())
493    }
494
495    #[test]
496    fn test_export_csv() -> Result<()> {
497        let mut animator = GradientFlowAnimator::new(0);
498        animator.record_step(0, &simple_grads(&["encoder"], 0.2));
499        animator.record_step(1, &simple_grads(&["encoder"], 0.18));
500
501        let path = temp_dir().join(format!("grad_anim_{}.csv", uuid::Uuid::new_v4()));
502        animator.export_csv(&path)?;
503        assert!(path.exists());
504        let content = std::fs::read_to_string(&path)?;
505        // Header + 2 data rows
506        assert!(content.lines().count() >= 3);
507        assert!(content.contains("step,layer,mean_abs_grad"));
508        Ok(())
509    }
510
511    #[test]
512    fn test_to_ascii_animation_empty() {
513        let animator = GradientFlowAnimator::new(0);
514        let out = animator.to_ascii_animation();
515        assert!(out.contains("no gradient frames"));
516    }
517
518    #[test]
519    fn test_to_ascii_animation_nonempty() {
520        let mut animator = GradientFlowAnimator::new(0);
521        for i in 0..5u64 {
522            animator.record_step(i, &simple_grads(&["embed", "attn"], 0.1 * i as f64));
523        }
524        let out = animator.to_ascii_animation();
525        assert!(out.contains("embed"));
526        assert!(out.contains("attn"));
527    }
528
529    #[test]
530    fn test_summary_report_healthy() {
531        let mut animator = GradientFlowAnimator::new(0);
532        for i in 0..5u64 {
533            animator.record_step(i, &simple_grads(&["layer"], 0.1));
534        }
535        let summary = animator.summary_report();
536        assert_eq!(summary.overall_health, GradientHealth::Healthy);
537        assert!(summary.layers_with_vanishing_grads.is_empty());
538        assert!(summary.layers_with_exploding_grads.is_empty());
539    }
540}
trustformers_debug/visualization/gradient_animation.rs

trustformers_debug/visualization/
gradient_animation.rs
