Neural Dimensionality Tracker (NDT) is a production-ready Python library for high-frequency monitoring of neural network representational dimensionality during training, enabling researchers and practitioners to track how internal representations evolve, detect discrete phase transitions (jumps), and gain mechanistic insights into the learning dynamics of deep neural networks across architectures (MLPs, CNNs, Transformers, Vision Transformers).
- Minimal Intrusion: Add dimensionality tracking to any PyTorch model with just 3 lines of code
- Architecture-Agnostic: Automatic support for MLPs, CNNs, Transformers, and Vision Transformers
- Multiple Metrics: Track 4 complementary dimensionality measures
- Jump Detection: Automatically identify phase transitions during training
- Rich Visualization: Built-in plotting with Matplotlib and interactive Plotly dashboards
- Flexible Export: Save results as CSV, JSON, or HDF5
- Production-Ready: Fully typed, tested (>90% coverage), and documented
pip install ndtrackerimport torch.nn as nn
from ndt import HighFrequencyTracker
# Your model
model = nn.Sequential(
nn.Linear(784, 512), nn.ReLU(),
nn.Linear(512, 256), nn.ReLU(),
nn.Linear(256, 10)
)
# Create tracker
tracker = HighFrequencyTracker(model, sampling_frequency=10)
# Training loop
for step, (x, y) in enumerate(dataloader):
output = model(x)
loss = criterion(output, y)
loss.backward()
optimizer.step()
tracker.log(step, loss.item()) # One line!
# Analyze
results = tracker.get_results()
from ndt import plot_phases
plot_phases(results, metric="stable_rank")- Quick Start: See above for 3-line integration
- Examples: examples/ contains complete working examples:
01_quickstart_mnist.py- Basic MLP on MNIST02_cnn_cifar10.py- CNN on CIFAR-1003_reproduce_tds_experiment.py- Reproduce TDS article experiment
- Full API Documentation: docs/
- Installation Guide: INSTALL.md
- Contributing: CONTRIBUTING.md
The repository includes a complete reproduction of the experiment described in the Towards Data Science article "I Measured Neural Network Training Every 5 Steps for 10,000 Iterations":
python examples/03_reproduce_tds_experiment.pyThis script uses the exact specifications from the article:
- Architecture: 784-256-128-10 (3-layer MLP)
- Dataset: MNIST (60k train/10k test)
- Optimizer: Adam (β1=0.9, β2=0.999)
- Learning rate: 0.001, batch size: 64
- Training: 8000 steps with measurements every 5 steps
- Expected results: 3 distinct phases (collapse, expansion, stabilization)
If you use Neural Dimensionality Tracker in your research, please cite:
@software{marin2024ndt,
author = {Marín, Javier},
title = {Neural Dimensionality Tracker: High-Frequency Monitoring of Neural Network Training Dynamics},
year = {2024},
publisher = {GitHub},
url = {https://github.com/Javihaus/ndt},
version = {0.1.0}
}Associated article:
@article{marin2025measuring,
author = {Marín, Javier},
title = {I Measured Neural Network Training Every 5 Steps for 10,000 Iterations: What High-Resolution Training Dynamics Taught Me About Feature Formation},
journal = {Towards Data Science},
year = {2025},
month = {November},
url = {https://towardsdatascience.com/}
}MIT License - see LICENSE file for details
Javier Marín
- LinkedIn: linkedin.com/in/jmarin
- Twitter: @javihaus
- Email: javier@jmarin.info
This work builds on research by:
- Achille, A., Rovere, M., & Soatto, S. (2019). Critical learning periods in deep networks. In International Conference on Learning Representations (ICLR). https://openreview.net/forum?id=BkeStsCcKQ
- Ansuini, A., Laio, A., Macke, J. H., & Zoccolan, D. (2019). Intrinsic dimension of data representations in deep neural networks. In Advances in Neural Information Processing Systems (Vol. 32, pp. 6109-6119). https://proceedings.neurips.cc/paper/2019/hash/cfcce0621b49c983991ead4c3d4d3b6b-Abstract.html
- Yang, J., Zhao, Y., & Zhu, Q. (2024). ε-rank and the staircase phenomenon: New insights into neural network training dynamics. arXiv preprint arXiv:2412.05144. https://arxiv.org/abs/2412.05144