Learning Computational Structure for Sample‑Efficient Prediction

Neurons get learned positions and connect only when close, so the network learns its own wiring.

Hangyu Zhou1, Constantine Dovrolis2, Steve Mussmann1

1Georgia Institute of Technology   2The Cyprus Institute

Paper Interactive demo Codecoming soon
LeWi networks learning six targets in turn: a depth-4 tree, a reuse DAG, a sampled DAG, flat modules, a high-arity node and a depth-5 tree. Each point is a neuron at its learned 3-D position and each line an active connection; colors show which piece of the target each neuron is associated with.

Many functions are built from small pieces

Learning the small pieces takes far less data than learning the whole function at once.

100,000,000cells for one function of 8 inputs, at 10 levels each
700cells for 7 pieces of 2 inputs

CSFRBench: the pieces are known, but hidden

A benchmark of targets built from small pieces on a graph we know. The learner sees only input–output examples; the graph is used only to score it.

TargetSmall pieces on a known graph, kept from the learner.
ExamplesRandom inputs and the output. This is all the learner sees.
LearnerTrained on the examples, then scored against the graph.

Standard networks miss most of the gain

A network given the target’s graph learns much faster. Dense and sparse baselines stay far behind. LeWi is not given the graph.

Let the neurons find their partners

  1. the learned position of neuron
  2. positive when and are within
  3. soft gate times binary mask
  4. the loss pulls and pushes
  5. crossing adds or removes a connection

Fewer samples when the target has structure

LeWi helps most when small pieces are stacked or reused.

The learned wiring aligns with the target

Neurons tied to the same piece of the target end up close together. Colors are assigned after training.

Each piece can be read outlinear probe R², best hidden layer, depth-4 tree
The network uses the piecesintervention R², best hidden layer, depth-4 tree
The wiring matches the target graphedge F1 of the learned module graph

It also helps world-model prediction

Over 16-step rollouts, LeWi beats a dense network on 16 of 17 tasks.

BibTeX

@misc{zhou2026lewi,
  title  = {Learning Computational Structure for Sample-Efficient Prediction},
  author = {Zhou, Hangyu and Dovrolis, Constantine and Mussmann, Steve},
  year   = {2026},
  url    = {https://zhou-hangyu.github.io/lewi/}
}