Exact Network Surgery and Reactive Computational Graphs in Julia with NeuroDSL

Specific Domains Machine Learning
, ,
1

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.

fig_f3_continuity
fig_f3_continuity1950×1200 337 KB

fig_e8_warmstart
fig_e8_warmstart750×500 35.2 KB

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

2

Hi, can your package be used with KeemenaLM.jl?

3

Thanks, you’re seemingly doing great work. I want to understand this more, so far scanning the papers.

Conclusion in paper ends with:

Together, these results turn “growing a network mid-training” from a fragile engineering trick into an operation with a specification, a proof, and a passing test suite.

At least I noticed in section 7.1:

A 4-block Llama-style model (d=128, 4 heads) is evaluated

While his GPT2 style, both are LLMs, seem similar enough (and both use Flux.jl, maybe a helpful thing here), but I don’t want to state confidently either way. Maybe @mantzaris would know. He has more heads, but I doubt that’s a conflict, since the above was only an evaluation, would apply to any (such) model with a different config (hyperparameters). I think “4-block” might mean 4 layers versus 24 there:

Section 5.3 (its title “.. Where The Speedup Evaporates”) is also intriguing.

From the other paper abstract:

enabling peak GPU memory reductions of 2-6x compared to PyTorch. .. proven numerically exact against full backpropagation on every topology we tested

I did see Net2Net in your paper, it’s also very intriguing:

4

@Palli Thanks for taking the time to read through the papers so carefully! You’ve highlighted some of the exact intersections between theory and systems engineering that motivated this project.

To address your points:

1. Compatibility with KeemenaLM.jl (and Flux.jl models)

You are completely right that conceptually, a GPT-2 style model and a LLaMA style model are very similar, and the surgical operations apply equally well to both. To clarify the “4-block” setup in the paper: yes, that refers to 4 transformer layers. Scaling this to KeemenaLM’s 24 layers and 16 heads just scales the graph size, but the \mathcal{O}(\vert{}\mathcal{V}\_s^+\vert{}) invalidation math holds perfectly.

However, direct “plug-and-play” with KeemenaLM.jl isn’t natively possible out of the box. KeemenaLM is built on Flux.jl, which relies on standard automatic differentiation, whereas NeuroDSL requires the model to be instantiated as a persistent, reactive DAG.

The workflow to bridge them would be: Instantiate the KeemenaLM topology in NeuroDSL, load the pre-trained weights from Flux, perform the “hot surgery” (e.g., adding a layer and warm-starting the AdamW state), and optionally export the weights back.

2. Section 5.3: “Where the Speedup Evaporates”

I’m glad you found this intriguing. In deep learning systems, there is no free lunch. While reactive invalidation is mathematically optimal—because it only recomputes the downstream cone of the insertion point—there is a constant bookkeeping overhead for the DAG mutation (measured around \sim 0.75 ms). If you mutate a layer that is too close to the input (meaning almost the whole network has to be recomputed anyway), or if the network is extremely shallow, that constant overhead overtakes the FLOPS you saved. Section 5.3 maps out that exact physical boundary.

3. The connection to Net2Net

Net2Net (Chen et al., 2015) is absolutely foundational here. It provided the brilliant mathematical blueprint for Identity Morphisms (how to pad weight matrices so the network function is preserved).

However, what NeuroDSL brings to the table is the systems infrastructure to actually execute Net2Net dynamically mid-training. In standard tape-based frameworks, expanding a network usually means pausing training, manually writing scripts to pad the weights, losing the optimizer state (or writing brittle code to pad the momentum buffers), and recompiling. Because every node in NeuroDSL natively owns its parameters and its optimizer state, we can execute Net2Net as a literal “hot surgery” and immediately warm-start the optimizer moments, drastically reducing the rejection rate of the new layers.

Thanks again for the great questions! Let me know if you want to dive deeper into the memory optimization side of the DAG.

5

Hi! Thanks so much for reaching out and for your interest in the project!

The short answer is: not directly “plug-and-play” out of the box, but yes, it is absolutely possible with a small conversion step.

Here is why: KeemenaLM.jl is built on top of Flux.jl, which uses a standard define-by-run approach. NeuroDSL handles things differently—it requires the model to be instantiated as a persistent, mutable DAG to allow for the “hot surgery” operations mid-training.

However, since KeemenaLM is fundamentally a standard transformer architecture (GPT-2 style), the workflow would be quite straightforward:

  1. Recreate the KeemenaLM architecture (layers, heads, dimensions) using NeuroDSL.
  2. Load the pre-trained weights from KeemenaLM into your new NeuroDSL graph.
  3. You are now ready to perform exact network surgery!

Here’s a sketch of what that conversion looks like in practice:

using NeuroDSL

# 1) Recreate the SAME architecture as your KeemenaLM/Flux model
#    (n_layers, dim, n_heads, hidden_dim, vocab_size must match exactly)
dev = NeuroDSL.Backend.CUDADevice()   # or CPUDevice()
g   = NeuroDSL.NeuroGraph(namespace=:keemena, device=dev)
ns  = :keemena

NeuroDSL.set!(g, :tokens, Float32.(your_token_ids); atom_type=NeuroDSL.Datom, namespace=ns)
x = NeuroDSL.Embedding(vocab_size, dim)(g, :tokens, :emb; namespace=ns)
# NB: LlamaModel has no built-in positional embedding or final norm --
# if your architecture uses either, add them the same way (another
# Embedding summed via an :add rule, and/or a LayerNorm) before/after
# the block below.

model = NeuroDSL.LlamaModel(n_layers, dim, n_heads, hidden_dim)
out   = model(g, x; namespace=ns)
logits = NeuroDSL.Linear(dim, vocab_size; bias=false)(g, out, :lmhead; namespace=ns)

# 2) Copy the trained weights from your Flux/KeemenaLM model into the
#    matching NeuroDSL symbols. NeuroDSL names every parameter
#    predictably:
#      layer_<i>_mha_q_W / _k_W / _v_W / _output_W   (attention projections)
#      layer_<i>_mlp_w1 / _mlp_w2 / _mlp_w3          (SwiGLU MLP)
#      layer_<i>_norm1_gamma / layer_<i>_norm2_gamma (RMSNorm scales)
#      emb_E, lmhead_W
#    Adapt the right-hand side below to however KeemenaLM actually
#    exposes its own parameters (this part is illustrative).

weights = Dict{Symbol,Array}(
    :emb_E    => keemena_embedding_matrix,   # (vocab_size, dim)
    :lmhead_W => keemena_lm_head_weight,     # (vocab_size, dim)
)
for i in 1:n_layers
    blk = keemena_model.blocks[i]            # <- adapt to KeemenaLM's real field names
    weights[Symbol(:layer_,i,:_mha_q_W)]      = blk.attn.wq
    weights[Symbol(:layer_,i,:_mha_k_W)]      = blk.attn.wk
    weights[Symbol(:layer_,i,:_mha_v_W)]      = blk.attn.wv
    weights[Symbol(:layer_,i,:_mha_output_W)] = blk.attn.wo
    weights[Symbol(:layer_,i,:_mlp_w1)]       = blk.mlp.w1
    weights[Symbol(:layer_,i,:_mlp_w2)]       = blk.mlp.w2
    weights[Symbol(:layer_,i,:_mlp_w3)]       = blk.mlp.w3
    weights[Symbol(:layer_,i,:_norm1_gamma)]  = blk.norm1.weight
    weights[Symbol(:layer_,i,:_norm2_gamma)]  = blk.norm2.weight
end

NeuroDSL.set_params!(g, ns, weights)   # handles Float16/BF16 -> Float32 + CPU/CUDA transfer

# 3) You're live -- insert_block!, patch_node!, greedy_patch_search!, etc.
#    all work on this graph exactly as on a NeuroDSL-native model.
NeuroDSL.demand!(g, logits; namespace=ns)

A couple of honest caveats: LlamaModel doesn’t include a positional embedding or a final norm by default — add those yourself the same way if KeemenaLM’s architecture has them. set_params! already handles the Float16/BFloat16 → Float32 conversion and CPU/CUDA transfer for you, so no manual casting needed on your end.

Happy to help debug the actual field-name mapping if you run into shape mismatches once you try it against KeemenaLM’s real internals!

6

If a Lux.jl backend was used in KeemenaLM.jl would it be possible to have them work together? That has not be implemented but the design does allow for different backends to be used if implemented.

7

Thank you very much for the detailed reply and for introducing me to NeuroDSL. It looks like a very interesting package!

Also, many thanks to @Palli for the additional information. While I can probably go head-to-head with @Palli in terms of decades of general computing experience, it seems in terms of ML and AI you guys are a bit more advanced.

I started learning Julia with AlphaZero.jl a few years ago, and I still believe that @jonathan-laurent’s toy problem examples are among the best in the entire Julia tour! As for neural networks, I have created one from scratch (about 2.5k lines) using Flux with some help from LLMs. It was Feed Forward Neural Network. I’ve also written some MCP servers and a RAG pipeline. In this field, the RAG pipeline is probably my best work so far. It’s a bit on the enterprise side, utilizing Cloudflare Workers and Oracle databases (Autonomous Database and MySQL HeatWave) and exposing several REST endpoints. I am also familiar with message brokers, with a particular focus on Redpanda and latency oriented POSIX inter-process communication. I usually write in Julia and C, and a bit of q.

I follow your point about the field-name mapping and shape mismatches! However, to be honest, I need a bit more time. There is a lot of new things for me on both the NeuroDSL and KeemenaLM sides.

Since we are on the topic, I would like to second @mantzaris’s question about Lux. I’ve been thinking about it as well. A few months (or perhaps years) ago, I was briefly in touch with @darsnack. My understanding at the time was that he suggested focusing on a new generation of ML, though he didn’t explicitly mention Lux back then. However, to be honest, I’m not sure if my understanding was correct.

P.S. BTW, if I may ask, did you catch the Spain vs. Argentina match?

8

Thank you for the question. As you noted, it hasn’t been implemented yet, but the separation of concerns in the graph engine’s design makes that kind of interoperability entirely feasible in principle.

NeuroDSL treats the computational graph as a persistent DAG where nodes own their cached values and operators are decoupled from the reactive invalidation logic. Because the backend primarily defines how individual node operators (or primitive layers) execute compute kernels (whether via CUDA.jl, native arrays, or a functional paradigm like Lux.jl), plugging in an alternative execution backend is a matter of mapping the graph’s primitive execution rules to that backend’s primitives, rather than redesigning the reactive engine itself.

The core challenge wouldn’t be structural or topological, but rather ensuring that the backend’s state-handling model (e.g., explicit parameters in Lux vs. mutable states) aligns cleanly with the dependency invalidation waves. It’s an open direction, and contributions or explorations in that direction are definitely welcome!

9

Really insightful, can you expand into more technical details on what would be needed for the state handling? On a high level I grasp what you are saying.

10

To get into the technical weeds: the friction (and the solution) comes from bridging a purely functional paradigm with a persistent stateful DAG.

In Lux.jl, a forward pass is strictly explicit: y, st_new = Lux.apply(layer, x, ps, st). The layer itself holds no state; parameters (ps) and states (st, like batch norm running means or PRNG keys) are passed in and returned as updated tuples.

In NeuroDSL, the computational graph is a persistent structure where each node owns its cached value and a valid boolean flag. When a mutation occurs, it triggers a reactive invalidation wave that is provably confined to the downstream cone.

To make them work seamlessly together, we would need a specific wrapper (e.g., a LuxNode operator) that acts as a translator between these two worlds. Here is what the technical implementation would require:

  1. Parameters as Reactive Roots: We cannot hide Lux’s ps inside a static struct within the node. Each parameter tensor from Lux must become its own source node in the NeuroDSL graph. This is crucial for exact surgery: if you mutate a specific Lux weight matrix, NeuroDSL must track that specific parameter node’s mutation to propagate the invalidation strictly to its downstream cone.

  2. State Threading (st): When Lux.apply returns st_new, the NeuroDSL wrapper must cache this new state locally. The technical trick here is managing the invalidation logic: updating st (like updating a BatchNorm running average during training) shouldn’t blindly trigger a reactive invalidation of the downstream graph for the current pass, but must be cached and ready for the next forward pass.

  3. The demand! Hook: The demand-driven evaluation in NeuroDSL would need to be overloaded for this LuxNode. When the engine calls demand! on an invalid Lux node, the node would:

    • Gather the cached inputs from its upstream input nodes.

    • Gather the cached ps from its upstream parameter nodes.

    • Run the explicit Lux.apply(layer, inputs, ps, st).

    • Cache the resulting y, update its internal st, and flip its flag to valid.

Essentially, you would “unroll” the explicit Lux parameters into NeuroDSL’s reactive dependency tracking, while keeping the Lux layer definitions as the underlying math engines. The design of NeuroDSL natively supports this kind of custom node definition, it just requires carefully wiring Lux’s explicit (ps, st) tuple into the topological sort!

11

Are you planning to update NeuroDSL in the near future by adding LuxNode wrapper to make it fully compatible with Lux, or is that currently rather outside your scope of interest?

12

To be completely transparent, while the technical blueprint for a LuxNode wrapper is clear, implementing it is not on my immediate short-term roadmap.

As an independent researcher, my current bandwidth is heavily dedicated to leveraging NeuroDSL’s reactive engine for fundamental research in mechanistic interpretability. In fact, I am currently finalizing a new paper (One Conditional, Two Scalpels) where this exact network surgery is crucial for studying the subset-dependent polarity of weight-space ablations. My primary focus right now is using the tool to explore these mathematical and structural behaviors in LLMs.

That being said, making NeuroDSL interoperable with the broader Julia ecosystem (like Lux) is absolutely within my scope of interest for the long-term evolution of the framework.

Since the project is fully open-source, I would warmly welcome community contributions. If you or anyone else is interested in experimenting with the LuxNode implementation we discussed, I would be more than happy to review PRs and provide detailed guidance on wiring Lux’s explicit state handling into the reactive topological sort!

13

Thank you for your detailed responses. I initially reached out to learn more about your project and to explore a potential collaboration as it seems that KeemenaLMis aspiring to be a very practical implementation of the subject matter.

While I find your project highly compelling, I am unable to commit the necessary time and effort at this stage without compromising my current priorities. Consequently, I will be looking forward to the release of the final implementation, hopefully in not so distant future.

Please be assured that this is solely due to current time constraints and not a lack of interest in contributing to your work. I trust this clarifies my reasoning and I hope this explanation provides the necessary context.

14

Thank you ! I completely understand ,balancing time constraints and current priorities is always the biggest challenge in both open-source and research.
I really appreciate your initial interest in the project and the thoughtful discussion we had. I will keep pushing forward with the core implementation of the reactive engine and the associated research in the meantime.
The door remains entirely open whenever your schedule frees up in the future.

15

Thank you, wishing you all the best, with hopes for future opportunities to connect!