Technical Summary: Learning Compositional Latent Structure with Vector Networks

Problem Statement

Modern deep networks typically rely on a shared dense parameter substrate to implement diverse computations. Through error backpropagation, distinct features, skills, and behaviors become encoded in overlapping directions of the same weight matrices, a phenomenon the authors term weight entanglement. In this regime, improvements for one behavior often inadvertently affect others because gradients are aggregated at the output and propagated through all layers. Consequently, standard networks lack "addressable responsibility" in weight space: they do not naturally store learned components as persistent, reusable objects that can be isolated, recombined, or updated independently. While post-hoc interpretability methods can sometimes identify functional circuits, the architecture itself does not enforce a decomposition into reusable units. This limitation hinders compositional generalization, where a system must recombine familiar factors into novel configurations without retraining. Existing approaches like Mixture-of-Experts or object-centric models introduce routing or modularity but often operate via feedforward routing over dense modules or do not make per-sample responsibility explicit in the weight space itself.

Methodology: The Vector Network (VN)

The Vector Network (VN) architecture addresses these limitations by replacing fixed, dense weight matrices with a library of reusable rank-1 weight atoms. The core mechanism involves a two-phase process: fast inference to compose input-specific transformations, and slow learning to update the underlying dictionary.

1. Architecture and Rank-1 Atoms

Each layer $l$ in a VN is parameterized not by a single matrix $W$, but by a dictionary of $K_l$ rank-1 weight atoms. Each atom $i$ is defined by a pair of learned code vectors: an input vector $s_{l,i}$ and an output vector $u_{l,i}$. Together, they form a rank-1 operator:
$$A^{(l,i)} = u_{l,i} s_{l,i}^\top$$
For any given input $x_l$, the network infers a sparse set of coefficients $g_l$ (where $g_{l,i}$ scales the contribution of atom $i$). The effective, input-specific weight matrix $W^*_l(x_l)$ is synthesized as a linear combination of these atoms:
$$W^*_l(x_l) = \sum_{i=1}^{K_l} g^*_{l,i}(x_l) A^{(l,i)}$$
This allows the network to compose a tailored transformation for each sample using a small subset of the shared dictionary, rather than applying a monolithic fixed matrix.

2. Fast Inference via Energy Minimization

For each sample, the VN performs fast inference to determine the sparse coefficients $g_l$. This is achieved by minimizing a layer-local energy function $E_l$ that balances three terms:

1. Input Reconstruction: Minimizing the error between the input $x_l$ and the reconstruction $S_l g_l$ (where $S_l$ contains the input-side code vectors).
2. Top-Down Consistency: Minimizing the error between the top-down target $h_{l,target}$ (provided by the layer above) and the interface message $U_l g_l$ (where $U_l$ contains the output-side code vectors).
3. Sparsity: An $\ell_1$ penalty $\lambda_l \|g_l\|_1$ to enforce a sparse active set.

The inference proceeds via iterative upward and downward sweeps (resembling an encoder-decoder settle) using proximal-gradient steps (ISTA/FISTA style). The process converges to a sparse fixed point where only a few atoms are active. The top layer may additionally enforce a hard top-$k$ constraint.

3. Slow Learning: Atomic-Hebb Updates

Once fast inference converges, slow learning updates the weight atoms. Crucially, updates are applied only to the atoms with non-zero coefficients ($g^*_{l,i} \neq 0$). The update rule, termed Atomic-Hebb, is a local residual update scaled by the inferred responsibility:
$$S_l \leftarrow S_l + \rho^S_l r^x_l (g^*_l)^\top$$
$$U_l \leftarrow U_l + \rho^U_l r^h_l (g^*_l)^\top$$
where $r^x_l$ and $r^h_l$ are the reconstruction and top-down residuals, respectively. Theoretical analysis (via Danskin's Theorem) proves that these local updates are exactly equivalent to the implicit gradients of the global free energy, despite the inference process being treated as a black box. This ensures that credit assignment is localized to the specific weight atoms responsible for the current explanation, preventing the diffuse weight entanglement seen in standard backpropagation.

Key Contributions

1. Rank-1 Weight Atoms as Building Blocks: The authors propose a hierarchical recurrent architecture where layers are composed of reusable rank-1 atoms, making network functions addressable, inspectable, and composable.
2. Inference-Defined Effective Weights: The effective weight matrix for a sample is not fixed but inferred via energy minimization over sparse coefficients. This separates the fast selection of components from the slow learning of the dictionary.
3. Atomic-Hebb Learning Rule: A learning rule that updates only the selected weight atoms using local residual signals scaled by inferred coefficients. This limits the overlap between unrelated updates and preserves the reusability of learned components.
4. Structural Limitation of Weight-Update Overlap: By restricting updates to the sparse active set, the architecture structurally limits weight entanglement, ensuring that learning one behavior does not necessarily degrade others.

Experimental Results

The authors evaluate VN on four benchmarks spanning 1D signals, 2D spatial decoding, N-body dynamics, and compositional MNIST. In all cases, models are trained in-distribution (ID) and tested on out-of-distribution (OOD) cases involving novel combinations of familiar factors.

• Spatial Composition: In a task requiring reconstruction of Gaussian bumps at unseen positions, standard CNNs and VAEs showed significant OOD error. VN maintained uniformly low error across the spatial domain, outperforming both dense baselines and hand-factorized decoders.
• Function Composition: On 1D signal reconstruction and masked prediction tasks involving sums and nested compositions of sine, cosine, polynomial, and Gaussian functions, VN demonstrated superior stability. While strong sequence models (Transformers, Mamba) performed well on ID and simple OOD reconstruction, they degraded significantly on masked prediction tasks requiring the recovery of reusable structure. VN achieved the clearest advantage in these stricter tests.
• N-Body Dynamics: In a physics benchmark involving mixed force fields (gravity, drag, Lorentz, spring) and varying body counts, VN maintained low error under complex OOD shifts (multiple forces, reduced bodies). Baselines like MLPs, GNNs, and Transformers showed large error increases when forces were combined or body counts changed, whereas VN remained robust.
• Image Composition: On MNIST digit reconstruction and dynamic transformation (Rotation/Translation/Scaling) tasks, VN significantly outperformed autoencoders in OOD settings involving additive mixtures, max projections, and complex transformation compositions. The advantage was roughly an order of magnitude in the strongest OOD conditions.

Across all benchmarks, VN matched strong baselines on ID performance but achieved substantially lower OOD error (often by an order of magnitude) when familiar factors had to be recombined in novel ways.

Significance and Claims

The paper argues that Vector Networks make compositional generalization a structural property of the architecture and inference process, rather than a brittle byproduct of fitting many behaviors into a shared dense substrate.

• Explicit Weight-Space Responsibility: Unlike standard networks where responsibility is implicit and entangled, VN explicitly infers which weight atoms are responsible for a sample before updating them. This "fast-recombination / slow-storage" split localizes plasticity.
• Mechanism-Level Composition: The results suggest that VN learns reusable, mechanism-level components (e.g., specific force laws or spatial factors) that can be recombined to explain novel inputs, rather than merely interpolating within the training distribution.
• Efficiency and Scalability: The authors suggest that once a library of latent factors is learned, new inputs can be handled by recomposition, offering a natural route toward improved data efficiency, particularly in the "long tail" of unseen combinations.
• Limitations: The paper acknowledges that the current implementation relies on recurrent sparse inference, which is slower than a single feedforward pass. The authors note that this inference could potentially be approximated by learned feedforward updates in future work, but the current form prioritizes interpretability and modularity.

In summary, the Vector Network proposes a shift from state routing (as in Transformers) to responsibility in weight space, where computation first infers sparse responsibility over reusable weight components, and learning then acts selectively on that inferred structure. This approach aims to solve the problem of weight entanglement and enable robust generalization to novel compositions of familiar factors.