Learning Computational Structure for Sample-Efficient Prediction

Hangyu Zhou; Constantine Dovrolis; Stephen Mussmann

cs.LGOct 20, 2026

Learning Computational Structure for Sample-Efficient Prediction

Authors: Hangyu Zhou, Constantine Dovrolis, Stephen Mussmann

Organizations: Georgia Institute of Technology · University of Cyprus

Abstract

Statistical learning in high dimensions suffers from the curse of dimensionality. Fortunately, many functions of practical interest are compositionally sparse, meaning they admit a composition of local functions, each depending on only a few variables, allowing sample complexity to depend on local arity rather than ambient dimension. Yet realizing this benefit requires exploiting hidden compositional structure, which theory does not establish for practical gradient-based training. We study this gap with CSFRBench, a controlled benchmark that hides the target's compositional computation during training but retains it for diagnosis, and find that dense networks and sparsity-based baselines leave much of that benefit unrealized. Inspired by the spatial organization of biological neural circuits, we introduce Learned Wiring (LeWi): each neuron has a learned low-dimensional embedding, only nearby neurons connect, and embeddings and weights are trained jointly. LeWi improves sample efficiency on targets with hierarchical structure, intermediate reuse, or independent modules. Across 32 additional targets generated by independently sampled computation graphs, LeWi is the most sample-efficient method not given the target graph, with 27% lower mean learning-curve area than dense networks. Dense networks perform better on targets with a single high-arity function. LeWi's representations and connectivity align with the target computation. On 17 world-model prediction tasks spanning control, navigation, and manipulation, LeWi beats the dense baseline on 16 and wins 99/119 baseline comparisons. These results show that jointly learning connectivity and weights can exploit computational structure to improve sample efficiency.

Figures & tables

Figure 1: LeWi develops embedding geometry aligned with the target computation. Learned neuron embeddings across successive layers of LeWi trained on a depth-4 tree. Colors indicate post hoc associations with intermediate values in the target computation. Inputs unused by the target and hidden units without a target association are omitted. Watch the geometry emerge during training : https://zhou-hangyu.github.io/lewi/ .

Figure 2: LeWi learns connectivity through neuron embeddings. Schematic illustration of neuron positions and gradient-driven movement. Adjacent-layer neurons connect within a fixed radius

R

(dashed), with shading indicating the soft gate. Prediction gradients can attract or repel embeddings, changing connectivity as pairs cross the radius. The zooms track one layer 1 neuron; trained colors illustrate post hoc associations with target intermediates.

Figure 3: Computation graphs in CSFRBench. Light nodes denote active inputs, blue nodes local functions, and green nodes readouts. The figure shows the five fixed graph templates and one draw from the sampled depth-4 DAG distribution.

Figure 4: LeWi improves sample efficiency across computational structures. Test NMSE versus training-set size for six representative CSFRBench settings. Curves and bands show means and 95% Student-

t

confidence intervals over target instances. Triangles mark means outside the displayed range. Baselines are shown in gray. Complete results are in Appendix D .

Figure 5: LeWi develops stronger target-aligned representations. Left: held-out $R^{2}$ from linear probes predicting each target intermediate from each hidden layer. Right: held-out counterfactual $R^{2}$ under joint multi-variable Boundless DAS. Rows are grouped by target level. Results are for the depth-4 tree at $n=2^{14}$ , averaged over 16 target instances.

Figure 6: Input selection, wiring alignment, and mechanism ablations. (a) LeWi–Dense paired AULC differences with all inputs relevant or 24 additional inert inputs. Hollow squares show medians when divergent Dense runs place means outside the plotted range. (b) Edge F1 over all target edges (filled) and edges from target intermediates (hollow) for trained LeWi and controls at

n=2^{15}

. (c) LeWi–variant paired AULC differences for four mechanism ablations. Negative differences favor LeWi in (a) and (c). Error bars show 95% Student-

t

confidence intervals over target instances.

Figure 7: World-model prediction. Paired $\Delta\mathrm{AULC}=\mathrm{AULC}(\mathrm{LeWi})-\mathrm{AULC}(\mathrm{baseline})$ across 17 tasks and seven baselines. Negative values favor LeWi. LeWi wins 99 of 119 comparisons, ties 6, and loses 14. Hollow circles indicate 95% confidence intervals containing zero. Colors are clipped at $\pm 0.75$ . Tasks are grouped by environment type and ordered by state dimension. Appendix E.3 reports complete results, including the learned-neuron-position comparison.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Local functions	Structure	Instances	Purpose
Single local function	Add.–mult.	arity 2	16	local fitting and input selection
Flat modules	Add.–mult.	four disjoint pairs	16	parallel low-arity computation
Depth-4 tree	Add.–mult.	balanced tree	16	hierarchical composition
Reuse DAG	Add.–mult.	reused intermediates	16	composition with intermediate reuse
Sampled depth-4 DAGs	Add.–mult.	sampled topology	32	generalization across graph draws
Depth-4 tree	Fourier	balanced tree	16	independent local-function family

Appendix

Table 1: CSFRBench evaluation settings. The suite spans parallel, hierarchical, reuse-based, sampled, and high-arity computation, together with an independent local-function family and a depth/scale extension.

Figure 8: Sampled computation graphs. Eight distinct topologies drawn from the depth-4 graph distribution. Each graph has eight active inputs, seven arity-two local functions, and three nonlinear levels, while branching, intermediate reuse, and readout connectivity vary.

Figure 9: Local-function fitting capacity. Median held-out NMSE over 16 development draws as a function of MLP width and hidden depth. Each model receives exactly the local-function arguments. (a,b) Arity-two local-function families. (c) Arity-eight additive–multiplicative stress setting. (d) Arity-eight Fourier control.

Figure 10: GraphOracle connectivity masks. Black entries denote retained weights. (a) GraphOracle(0), GraphOracle(1), and GraphOracle(all) on the depth-4 tree. (b) GraphOracle(all) on the reuse DAG, flat modules, and high-arity node.

Method	Parameter	$32\times 4$ backbone	Depth-5 tree ( $64\times 8$ )
Dense	learning rate	$[0.00143,\,0.261]$	$[7.70\times 10^{-4},\,0.475]$
$\ell_{1}$	learning rate	$[1.57\times 10^{-6},\,0.349]$	$[2.72\times 10^{-4},\,0.457]$
	$\lambda$	$[10^{-7},\,670]$	$[10^{-7},\,6.28\times 10^{11}]$
$\ell_{2}$	learning rate	$[2.16\times 10^{-5},\,0.648]$	$[1.93\times 10^{-5},\,0.115]$
	$\lambda$	$[10^{-7},\,1.13\times 10^{3}]$	$[10^{-7},\,7.26\times 10^{3}]$
$L_{0}$	learning rate	$[2.67\times 10^{-4},\,0.0443]$	$[2.46\times 10^{-4},\,0.00839]$

Appendix

Table 2: Hyperparameter search spaces. Frozen ranges used in the final CSFRBench evaluation and reported appendix experiments. All searched numerical parameters are sampled log-uniformly. Search ranges are constructed on development targets before final evaluation. The depth-5 tree uses separately frozen ranges because it has a different backbone and training horizon. A dash indicates that an experiment is not evaluated on that backbone.

Figure 11: Complete CSFRBench learning curves. Held-out test NMSE versus training-set size for all evaluated methods and settings. Fixed settings are averaged over 16 target instances and sampled graphs over 32 graph/function instances. Bands show 95% Student-

t

confidence intervals for highlighted methods, and triangles mark means outside the displayed range.

Figure 12: Paired predictive comparisons across CSFRBench. Paired AULC differences between LeWi and every baseline or graph-informed reference. Negative values indicate lower AULC for LeWi. Points and intervals show means and 95% Student-

t

confidence intervals over target instances. Triangles mark means outside the displayed range, and hollow diamonds show the corresponding medians.

	Depth-4	Depth-3	Depth-4	Flat	Fourier	High-arity
Method	tree	reuse DAG	Fourier tree	modules	reuse DAG	node
LeWi	$0.508{\pm}0.062$	$0.179{\pm}0.042$	$0.035{\pm}0.013$	$0.008{\pm}0.001$	$0.019{\pm}0.006$	$0.085{\pm}0.011$
Dense	$0.633{\pm}0.059$	$0.271{\pm}0.056$	$0.071{\pm}0.023$	$0.018{\pm}0.002$	$0.043{\pm}0.009$	$0.066{\pm}0.008$
$\ell_{1}$	$0.626{\pm}0.061$	$0.268{\pm}0.061$	$0.044{\pm}0.020$	$0.013{\pm}0.002$	$0.027{\pm}0.007$	$0.066{\pm}0.008$
$\ell_{2}$	$0.593{\pm}0.061$	$0.220{\pm}0.056$	$0.046{\pm}0.019$	$0.015{\pm}0.002$	$0.027{\pm}0.005$	$0.048{\pm}0.006$
$L_{0}$	$0.569{\pm}0.065$	$0.232{\pm}0.060$	$0.056{\pm}0.020$	$0.027{\pm}0.004$	$0.038{\pm}0.008$	$0.089{\pm}0.010$

Appendix

Table 3: Mean AULC on the fixed CSFRBench settings. Values are means over 16 target instances with 95% Student-

t

confidence half-widths. AULC integrates test NMSE over

n=2^{9},\ldots,2^{15}

; lower is better.

	Depth-5 tree		Independently sampled depth-4 DAGs
Method	AULC	LeWi $-$ method	AULC	LeWi $-$ method
LeWi	$0.812{\pm}0.037$	–	$0.197{\pm}0.030$	–
Dense	$1.185{\pm}0.572$	$-0.374$ [ $-0.950$ , $+0.203$ ] (16/16)	$0.268{\pm}0.033$	$-0.072$ [ $-0.080$ , $-0.063$ ] (32/32)
$\ell_{1}$	$0.905{\pm}0.020$	$-0.093$ [ $-0.114$ , $-0.073$ ] (16/16)	$0.244{\pm}0.031$	$-0.047$ [ $-0.056$ , $-0.038$ ] (32/32)
$\ell_{2}$	$0.895{\pm}0.021$	$-0.084$ [ $-0.103$ , $-0.065$ ] (16/16)	$0.245{\pm}0.033$	$-0.049$ [ $-0.059$ , $-0.039$ ] (32/32)
$L_{0}$	$0.882{\pm}0.025$	$-0.070$ [ $-0.089$ , $-0.052$ ] (16/16)	$0.245{\pm}0.033$	$-0.048$ [ $-0.056$ , $-0.040$ ] (32/32)

Appendix

Table 4: AULC on the depth-5 tree and on independently sampled depth-4 DAGs. Mean AULC with 95% Student-

t

half-widths, and

\mathrm{AULC}(\mathrm{LeWi})-\mathrm{AULC}(\mathrm{method})

as mean [95% paired Student-

t

interval] over the number of instances on which LeWi has lower AULC; negative values favor LeWi. The depth-5 tree (16 target instances; width 64, eight hidden layers, 128,000 steps) integrates test NMSE over

n=2^{11},\ldots,2^{17}

, the 32 independently sampled depth-4 DAG/function instances over

n=2^{9},\ldots,2^{15}

; neither range matches Table 3 , so AULC values are comparable within a column only.

Figure 13: Representation diagnostics for additional learned baselines. Left: held-out $R^{2}$ from ridge probes predicting each target intermediate from each hidden layer. Right: held-out counterfactual $R^{2}$ under joint multi-variable Boundless DAS. The top row repeats Dense, LeWi, and GraphOracle(all) from Figure 5 , while the bottom row shows $L_{0}$ , RigL, and BIMT. Rows are grouped by target level and columns correspond to hidden layers. Results are for the depth-4 tree at $n=2^{14}$ , averaged over 16 target instances.

Figure 14: Representation diagnostics across training-set size and target structure. (a) Maximum held-out ridge-probe

R^{2}

on the depth-4 tree, reuse DAG, and Fourier tree. For each target intermediate, we take the maximum

R^{2}

across hidden layers and then average across target intermediates and target instances. (b) Mean held-out counterfactual

R^{2}

under joint multi-variable Boundless DAS on the same target families. Curves show Dense, LeWi, and GraphOracle(all), averaged over 16 target instances with 95% Student-

t

confidence intervals.

Figure 15: Complete CSFRBench learning curves with nuisance inputs. Held-out test NMSE versus training-set size with ambient dimension

d=32

. The depth-4 settings contain eight active and 24 inert inputs, while the depth-5 tree contains 16 active and 16 inert inputs. Fixed settings are averaged over 16 target instances and sampled graphs over 32 graph/function instances. Bands show 95% Student-

t

confidence intervals for highlighted methods, and triangles mark means outside the displayed range.

Figure 16: Paired predictive comparisons with nuisance inputs. Paired AULC differences between LeWi and every baseline or graph-informed reference at

d=32

. Negative values indicate lower AULC for LeWi. Points and intervals show means and 95% Student-

t

confidence intervals over target instances. Triangles mark means outside the displayed range, and hollow diamonds show the corresponding medians.

Figure 17: Representative target and extracted module graphs. Each column shows a target computation graph (top) and the module graph extracted from one trained LeWi network (bottom). Colors denote post hoc matches between modules and target intermediates. Black edges correspond to target dependencies, red target edges are missed, and the remaining edge styles indicate spurious connections, transitive shortcuts, or paths through unmatched modules. These examples illustrate the graph comparison summarized by edge precision, recall, and

F_{1}

. Each example is the network with upper-median edge

F_{1}

in its setting.

Setting	All edges			From intermediates			Assessable edges
	Precision	Recall	$F_{1}$	Precision	Recall	$F_{1}$	Recall	Share
Depth-4 tree	$0.677$	$0.579$	$0.623$	$0.503$	$0.375$	$0.425$	$0.731$	$0.800$
Reuse DAG	$0.886$	$0.695$	$0.777$	$0.829$	$0.493$	$0.613$	$0.772$	$0.915$
Fourier tree	$0.518$	$0.446$	$0.478$	$0.403$	$0.321$	$0.354$	$0.508$	$0.895$
Sampled depth-4 DAGs	$0.674$	$0.509$	$0.575$	$0.634$	$0.423$	$0.499$	$0.542$	$0.935$

Appendix

Table 5: Wiring alignment at

n=2^{15}

. Precision, recall, and

F_{1}

between the extracted module graph and the target computation graph, reported over all target edges and over edges originating at target intermediate nodes. Assessable recall restricts to target edges whose endpoints have matched modules.

Figure 18: Wiring alignment across training-set size. Edge F1 of the extracted module graph against the target computation graph as the number of training samples increases. Filled markers report F1 over all target edges; hollow markers report F1 for edges originating at target intermediate nodes. Curves show trained LeWi and the corresponding initialization and randomized controls, averaged over target instances with 95% Student-

t

confidence intervals.

Figure 19: Learned neuron embeddings across target structures. Examples from trained LeWi models on the depth-4 tree, reuse DAG, flat modules, high-arity node, and a sampled depth-4 DAG. Each row shows the target computation graph followed by the learned three-dimensional neuron embeddings across successive network-layer transitions. Lines denote active connections, and colors indicate post hoc associations with target intermediates. Additional embedding visualizations and their training trajectories are available in the interactive demo linked in Figure 1 .

Figure 20: Learned connection density. Active-edge fraction at the selected checkpoint across training-set size with (a) all inputs active and (b) nuisance inputs at

d=32

. The dotted line marks the initialization density of

0.61

. The depth-5 tree uses

d=16

in (a) and 16 active plus 16 nuisance inputs in (b); the other settings use 8 active inputs and, in (b), 24 nuisance inputs. Curves show means over target instances with 95% Student-

t

confidence intervals.

Figure 21: Sensitivity to LeWi’s geometric parameterization. Paired AULC differences from the adopted configuration when varying the embedding dimension

r

(left), connection threshold

R^{2}

with

\sigma_{e}

fixed (middle), and initialization scale

\sigma_{e}

with

R^{2}=2r\sigma_{e}^{2}

(right). Positive values indicate worse sample efficiency than the adopted configuration. Results are shown for the depth-4 tree and Fourier tree. Error bars show 95% Student-

t

confidence intervals over target instances. Dashed lines mark the adopted values

r=3

R^{2}=1.5

, and

\sigma_{e}=0.5

	Parameters		Runtime (s)
Backbone	Dense	LeWi	Dense	LeWi	Ratio
$32\times 4$	3,489	3,900	5.2	6.8	1.32
$64\times 8$	30,273	31,860	84.3	85.6	1.02

Appendix

Table 6: Computational cost of LeWi and Dense. Parameters and runtime (seconds per trained network) using optimized CPU implementations. Runtime includes training, validation, and testing, with both methods run concurrently on the same node (one thread per process).

Method	Parameter	Search range
Dense	learning rate	$[8.58\times 10^{-5},\,4.63\times 10^{-2}]$
$\ell_{1}$	learning rate	$[4.94\times 10^{-5},\,9.85\times 10^{-2}]$
	$\lambda$	$[1.66\times 10^{-9},\,1.66\times 10^{3}]$
$\ell_{2}$	learning rate	$[3.44\times 10^{-4},\,3.98\times 10^{-2}]$
	$\lambda$	$[5.22\times 10^{-7},\,9.06\times 10^{1}]$
$L_{0}$	learning rate	$[1.30\times 10^{-3},\,6.50\times 10^{-2}]$

Appendix

Table 7: Hyperparameter search spaces for the world-model experiments. All searched numerical parameters are sampled log-uniformly.

Figure 22: Complete world-model comparisons. Each cell shows the mean paired $\Delta\mathrm{AULC}=\mathrm{AULC}(\mathrm{LeWi})-\mathrm{AULC}(\mathrm{baseline})$ over 16 paired seeds; negative values favor LeWi. A dot marks a comparison whose 95% Student- $t$ confidence interval contains zero. The right margin reports LeWi’s win–tie–loss count against each baseline across the 17 tasks.

Figure 23: Complete world-model learning curves. Rollout-16 NMSE versus the number of training transitions for all 17 tasks and nine methods. Curves show arithmetic means over 16 paired final seeds with 95% Student-

t

confidence intervals; lower is better. Triangles indicate means outside the displayed range. Learned neuron positions are included here but omitted from the main-text heatmap because their larger errors would compress its color scale.

Data-Efficient Learning

Data-efficient learning develops methods that achieve accurate, robust models with fewer labeled examples, training samples, or interactions.

See all 243 papers

May 12, 2026

Understanding Sample Efficiency in Predictive Coding

Predictive Coding (PC) is an influential account of cortical learning. Much of recent work has focused on comparing PC to Backpropagation (BP) to find whether PC offers any advantages. Small scale experiments show that PC enables learning that is more sample efficient and effective in many contexts, though a thorough theoretical understanding of the phenomena remains elusive. To address this, we quantify the efficiency of learning in BP and PC through a metric called ``target alignment'', which measures how closely the change in the output of the network is aligned to the output prediction error. We then derive and empirically validate analytical expressions for target alignment in Deep Linear Networks. We show that learning in PC is more efficient than BP, which is especially pronounced in deep, narrow and pre-trained networks. We also derive exact conditions for guaranteed optimal target alignment in PC and validate our findings through experiments. We study full training trajectories of linear and non-linear models, and find the predicted benefits of PC persist in practice even when some assumptions are violated. Overall, this work provides a mechanistic understanding of the higher learning efficiency observed for PC over BP in previous works, and can guide how PC should be parametrised to learn most effectively.

Gaspard Oliviers, Elene Lominadze, Rafal Bogacz

Nuffield Department of Clinical Neurosciences, University of Oxford, United Kingdom · MRC Centre of Research Excellence in Restorative Neural Dynamics, United Kingdom

Backpropagation Predictive Coding Deep Linear Networks Neural Network Training Dynamics Data-Efficient Learning

Ask PDF

May 25, 2026

Learning Sparse Compositional Functions with Norm-Constrained Neural Networks

The ability of deep neural networks to learn hierarchical features is widely regarded as a key mechanism underlying their success in high-dimensional learning. Existing theory partially supports this view by establishing approximation rates based on parameter counts and sample complexity guarantees for compositional models without incurring the curse of dimensionality (CoD). To study overparameterized regimes, where the number of parameters exceeds the sample size, we develop a framework that measures complexity via the parameter norm. Within this approach, we establish approximation rates and excess risk bounds for learning sparse compositional functions whose compositional structure is represented by directed acyclic graphs (DAGs), using Frobenius norm-constrained deep neural networks. Our results have broad applicability since every function that is efficiently Turing computable admits sparse compositional representations. In particular, we cover a range of representative models, including multi-index models, binary tree structures, and general compositional architectures. The rates we derive show that deep networks can exploit the compositional structure of the target functions, effectively avoiding the CoD through hierarchical representations.

Shuo Huang, Lorenzo Fiorito, Lorenzo Rosasco +1

MaLGa, DIMA, Università degli Studi di Genova, Genoa, Italy · CBMM, Massachusetts Institute of Technology, Cambridge, MA, USA

Neural Network Generalization Hierarchical Representation Learning Neural Network Approximation Theory Compositional Representation Learning

Ask PDF

May 14, 2026

Learning with Shallow Neural Networks on Cluster-Structured Features

The success of deep learning in high-dimensional settings is often attributed to the presence of low-dimensional structure in real-world data. While standard theoretical models typically assume that this structure lies in the target function, projecting unstructured inputs onto a low-dimensional subspace, data such as images, text or genomic sequences exhibit strong spatial correlations within the input space itself. In this paper, we propose a tractable model to study how these correlations affect the sample complexity of learning with gradient descent on shallow neural networks. Specifically, we consider targets that depend on a small number of latent Boolean variables, and input features grouped into clusters and correlated with the latent variables. Under an identifiability assumption, we show that for a layerwise gradient-descent variant, the sample complexity scales with the number of hidden variables and, when the signal-to-noise ratio is sufficiently high, is independent of the input dimension, up to logarithmic terms. We empirically test our theoretical findings on both synthetic and real data.

Elisabetta Cornacchia, Laurent Massoulié

INRIA, DI/ENS, PSL.

Shallow Neural Networks Neural Network Generalization Sample Complexity Statistical Learning Theory

Ask PDF

Abstract

Data-Efficient Learning

Explore similar work

Understanding Sample Efficiency in Predictive Coding

Learning Sparse Compositional Functions with Norm-Constrained Neural Networks

Learning with Shallow Neural Networks on Cluster-Structured Features