Exact Network Surgery and Reactive Computational Graphs in Julia with NeuroDSL
Hi everyone,
I’d like to share some recent theoretical and systems results from NeuroDSL, a persistent, reactive computational graph framework for Deep Learning built entirely in Julia.
While tape-based and define-by-run frameworks dominate ML, they struggle with topological mutations mid-training: adding a layer usually means discarding compiled artifacts and manually re-associating the optimizer state. Julia allows us to take a different approach.
By treating the computational graph as a persistent, mutable DAG where nodes own both their values and their optimizer states, NeuroDSL achieves Exact Network Surgery.
Here are a few things this architecture uniquely enables, as detailed in my latest preprint:
-
\mathcal{O}(\vert{}\mathcal{V}\_s^+\vert{}) Invalidation : Mutating the graph triggers a reactive invalidation wave that recomputes exactly the downstream cone of the insertion point. The graft-plus-invalidation bookkeeping is constant (measured at ~0.75 ms) across insertion depths.
-
Bit-Exact Identity Morphisms : Inserting a gated residual block preserves the network function bit-exactly (verified preconditions, 0 mismatches out of 1600 logits on a Llama-style block) while ensuring branch gradients unlock immediately.
-
Counterfactual Optimizer Interventions : Because state is persistent per-node, we tracked down why networks initially “reject” newly grafted layers. It turns out to be a cold-start artifact of AdamW’s second moment (v). By doing an on-the-fly “warm-start” of the branch’s v to match its downstream consumer—a trivial operation in NeuroDSL but a nightmare in tape-based systems—we slowed the rejection rate by nearly an order of magnitude (6.9x to 11.6x).
The paper also formalizes the aggregate speedup of exhaustive sweeps using the Karamata index of the network’s cost-by-depth profile, and proves the exact interleaved and batched cost of persistent grafts.
I believe Julia’s multiple dispatch and meta-programming capabilities are uniquely suited for this kind of reactive Differentiable Programming. I’d love to get feedback from the autodiff and ML systems community here.
Interractive Graph surgery
Github repo
Article : Exact Network Surgery