Multiway Multislice PHATE: Visualizing Hidden Dynamics of RNNs through Training

TL;DR

This paper introduces MM-PHATE, a visualization method that reveals how RNN internal representations evolve during training across multiple dimensions, aiding understanding and development of better models.

Contribution

The paper presents MM-PHATE, a novel graph-based embedding technique for visualizing RNN hidden states across time, training epochs, and units, capturing dynamic changes during learning.

Findings

MM-PHATE preserves community structure in hidden states.

It reveals training-phase changes in representation geometry.

Aligns with probes and mutual information metrics.

Abstract

Recurrent neural networks (RNNs) are a widely used tool for sequential data analysis; however, they are still often seen as black boxes. Visualizing the internal dynamics of RNNs is a critical step toward understanding their functional principles and developing better architectures and optimization strategies. Prior studies typically emphasize network representations only after training, overlooking how those representations evolve during learning. Here, we present Multiway Multislice PHATE (MM-PHATE), a graph-based embedding method for visualizing the evolution of RNN hidden states across the multiple dimensions spanned by RNNs: time, training epoch, and units. Across controlled synthetic benchmarks and real RNN applications, MM-PHATE preserves hidden-representation community structure among units and reveals training-phase changes in representation geometry. In controlled synthetic systems spanning multiple bifurcation families and smooth state-space warps, MM-PHATE recovers qualitative dynamical progression while distinguishing family-level differences. In task-trained RNNs, the embedding identifies information-processing and compression-related phases during training, and time-resolved geometric and entropy-based summaries align with linear probes, time-step ablations, and label--state mutual information. These results show that MM-PHATE provides an intuitive and comprehensive way to inspect RNN hidden dynamics across training and to better understand how model architecture and learning dynamics relate to performance.