$ show astroautomata/SymTorch
SymTorch is a Python library for deep learning interpretability via symbolic regression. It wraps PyTorch neural network components (layers, submodules) or arbitrary callable functions (e.g., LLM inference, scikit-learn models) and uses the PySR symbolic regression engine (Julia-backed) to discover human-readable mathematical equations that approximate the wrapped component's input-output behavior. The library provides a unified SymbolicModel class that can operate at the layer level (using forward hooks to capture intermediate activations) or in a model-agnostic mode for any black-box function. It also supports local interpretability via SLIME (synthetic neighborhood sampling with weighted regression) and progressive dimensionality reduction through configurable pruning schedules (exponential, linear, cosine decay). The primary workflow involves wrapping a model, optionally training with pruning, calling distill() to run symbolic regression, inspecting discovered equations via show_symbolic_expression(), and optionally switching the forward pass to use the symbolic equation instead of the original block via switch_to_symbolic(). The architecture is monolithic: the SymbolicModel class (~2300 lines) encapsulates all functionality including caching, PySR parameter management, forward hook handling, SLIME sampling, pruning schedules, and serialization via dill. While this centralization simplifies the public API, it creates a single point of maintenance and limits modular reuse. The library relies on several heavy external dependencies: PySR (which requires Julia), PyTorch, SymPy, and dill. Documentation includes comprehensive API reference, Jupyter notebook demos (PINNs, GNNs, LLM arithmetic probing), and case studies demonstrating symbolic extraction from physics-informed neural networks and small language models. The codebase demonstrates solid engineering with good docstrings, extensive caching to avoid redundant computations, and thoughtful design decisions like configurable pruning schedules and SLIME parameters. However, there are signs of academic development: future-dated version numbers (1.1.1 docs vs 1.0.1.post1 setup.py), local path references in demos, and incomplete decomposition (TODO comments about breaking up the class). The testing suite exists but heavily mocks PySR, which may miss integration issues. Overall, SymTorch is a functional, research-oriented interpretability toolkit with clear value for extracting symbolic equations from trained models, but its monolithic structure and external dependencies limit its reuse as a standalone component.
The LLM identified 3 files worth highlighting for reuse:
The SLIME sampling logic (methods `_apply_slime_sampling` and related parameters) generates synthetic neighborhood points around a query point using Gaussian perturbations, selects nearest neighbors via a configurable metric, and computes distance-based weights. This is a self-contained, model-agnostic local explanation pattern that can be extracted as a standalone utility for any black-box model to produce weighted local samples for surrogate model fitting.
The pruning schedule generator (`_set_pruning_schedule`) implements three decay strategies (exponential, linear, cosine) to progressively reduce the number of active output dimensions during training. It takes initial/target dimensions, total steps, and decay parameters, returning a schedule of target dimensions per step. This function is isolated and reusable in any neural network pruning or dynamic architecture framework without dependence on the rest of SymbolicModel.
The caching mechanism (`_check_cache_hit`, `clear_cache`, `_cache` dict) stores input-output pairs from forward passes to avoid redundant computations across multiple `distill` calls. The cache key is based on a hash of inputs and parent model identity, ensuring freshness when the underlying model changes. This pattern is a simple, effective optimization for any iterative model analysis workflow and can be extracted as a generic caching decorator or context manager.