Pioneers Insight Method Research Author
Neural Networks Are Elastic Origami! [Prof. Randall Balestriero]
Back to Episodes

Neural Networks Are Elastic Origami! [Prof. Randall Balestriero]

Summary

  • Balestriero’s core claim is that modern deep networks are continuous piecewise-affine splines: “elastic origami” that partitions input space, stretches each region, then uses a linear output map to separate classes. The geometry exposes where capacity, brittleness, bias, and extrapolation reside, creating practical design levers that ordinary loss and accuracy curves hide.
  • Across broader computer-vision settings, including CIFAR, clean train and test accuracy evolved and plateaued at roughly the same time, but adversarial robustness appeared only after much longer ordinary training. Regions migrated away from samples and accumulated at decision boundaries, leaving wider, nearly affine neighborhoods around data. No adversarial examples were used, so Balestriero argues the resulting protection should be less attack-specific—while hedging that adversarial grokking occurred in “most of the settings that we found.”
  • Under a fixed FLOP budget, a smaller model trained longer might acquire robustness that a maximal model stopped at the clean-accuracy plateau never reaches. This is the episode’s “smart FLOPs” call, but standard tools can oppose it: batch normalization concentrates regions near examples, while stronger weight decay favors smooth near-zero solutions and may delay or prevent the sparse endpoint.
  • Pixel-space reconstruction inherits a dataset’s spectral bias because mean-squared error supplies its strongest gradients through high-energy, low-frequency structure, while class-relevant detail often arrives later in lower-energy, high-frequency features. Designed masking or noise can rebalance learning but requires task knowledge; latent-space objectives more easily discard irrelevant pixels, explaining why attractive reconstructions can coexist with mediocre perception features.
  • Seven geometric features per LLM MLP block—about 500 even for a 70B model—were enough to cluster prompts by domain and toxicity and support a small linear detector. On OmniToxic, the host cites roughly 73.5% area under the curve for a popular detector versus 99.18% for the geometry probe on Llama 2 7B, with latency adjustable by reading only the first few layers.
  • The discussion frames alignment risk as a possible complexity ceiling: longer, semantically related context lowers attention sparsity, raises a prompt’s inferred intrinsic dimension, and can move it beyond the space covered by RLHF. Balestriero reports jailbreaks on Llama 2 and an additional ChatGPT experiment; natural related sentences can outperform random padding and are harder to filter. Scalable alignment therefore needs architectural or geometric guarantees, not only more preference examples.

Deep dive

1. A deep network is “elastic origami,” not an opaque function

  • Balestriero’s 2018-era work was not inventing splines, which he says were probably used since the 1980s or earlier, but applying their theory to contemporary deep networks. Compositions of affine operations—dense layers or convolutions—with ReLU, max pooling, or similar piecewise-linear nonlinearities produce one continuous piecewise-affine input-output map.

  • The resulting geometry cuts a high-dimensional image or token space into convex regions. Inside each region the network performs only an affine mapping; the nonlinear behavior occurs when an input crosses a boundary. His preferred picture is “origami,” except the space also stretches: “I would say it’s elastic origami.”

  • The host’s challenge—worth keeping—is that local templates or locality-sensitive hashing do not resemble reasoning. Balestriero agrees that each region looks like template matching, but shared parameters couple distant regions: learning from one example alters behavior where no data exists, allowing hidden rules to transfer to “another example that you did not see during training.”

  • He says the viewpoint is not limited to MLPs or vision: a convolution can be treated as an MLP with a circulant, constrained matrix, and the affine-plus-nonlinearity structure extends across current architectures and modalities.

2. Partition placement buys more than richer local polynomials

  • Classical spline approximation supplies the key design lesson: if forced to choose between increasing each region’s polynomial degree and positioning simple regions around the data, adaptive placement is much better. Piecewise-affine functions can therefore be “optimal if you can position those regions well,” despite their minimal within-region complexity.

  • Training simultaneously learns the affine maps and the partition boundaries because both are tied to the same weights. Regions concentrate around training data and along extrapolation patterns imposed by the architecture; remote areas outside both receive far fewer regions. Smaller local regions mean more precise approximation.

  • This distribution makes geometry sample-specific. If one demographic receives fewer regions and thus less local expressivity than another, Balestriero says the network can be described as algorithmically biased there. Partition statistics replace the vague question “what does this network do?” with “what does it do in this neighborhood?”

3. Adversarial robustness groks after clean accuracy looks finished

  • Conventional grokking describes delayed generalization: the training metric reaches a plateau while test performance can remain near random or only slightly above it, then rises after training continues far beyond the normal stopping point. Gradient information keeps rearranging weights even though the visible training metric appears finished.

  • Balestriero and his co-author tested broader settings—including CIFAR and other computer-vision tasks, convolutional networks, and ResNets—rather than only specially chosen tasks, models, and initializations. They did not find delayed clean generalization: clean train and test accuracy evolved at roughly the same rate and plateaued at roughly the same time.

  • The delayed property instead appeared under adversarial perturbation. A white-box attack uses network gradients to find a tiny input change that causes misclassification; a network with perhaps 100% clean training accuracy can fall to random guessing under changes “that you cannot see by eye.”

  • Continuing ordinary training eventually improved adversarial test accuracy, despite using no adversarial training. Balestriero calls this “adversarial grokking” and reports it in most settings examined, not unconditionally everywhere. Its timing is invisible if monitoring stops at the clean-accuracy plateau.

4. Late training moves capacity from samples to decision boundaries

  • Early training aggressively packs regions around individual samples. The model is already capable of extrapolation, but geometrically it resembles memorization: many parameters resolve the observed points, while complexity remains spread relatively uniformly through the surrounding space.

  • During much longer training, ongoing gradient information initiates a local “decomplexification.” Regions migrate away from training and test points and become densely compressed near decision boundaries. Neighborhoods around samples consequently widen, making the model affine over a larger radius.

  • That redistribution explains robustness: a small perturbation is less likely to cross a nonlinear boundary. Capacity is retained where curvature is actually necessary, at the transition between classes, rather than spent repeatedly encoding stable neighborhoods. Balestriero says the mapping tends toward piecewise-constant behavior in the theoretical limit.

  • The host notes that the late geometry resembles a topographic or Voronoi map, with many boundaries “squashed together” between classes. Whether one calls the process complexification or decomplexification depends on location: sample neighborhoods simplify while the class boundary receives a highly concentrated representation budget.

5. Sparsity and double descent describe the same geometric migration

  • Balestriero connects the endpoint to pruning, rank collapse in the parameters, neural collapse, and sparse networks. He says there is a one-to-one correspondence between these effects, different regularizers, and their impact on partition geometry.

  • Refined pruning can simplify a partition by deleting units responsible for unnecessary boundaries near samples, while retaining boundaries that resolve the decision surface. His lottery-ticket interpretation is geometric: iterative magnitude pruning may forcibly replace the early, uniformly complicated solution with something closer to the late-training sparse solution.

  • Robustness can improve even after most weights disappear because the removed capacity was positioned where the task did not need it. The remaining geometry concentrates representation capacity near the decision boundary.

  • Local complexity also produces a double-descent-like curve. Region counts rise around points as clean accuracy improves, peaking near the clean plateau with little robustness; during the second descent, regions move away, local radii expand, and adversarial robustness eventually appears.

6. Standard regularization can block the robustness endpoint

  • Regularization is not neutral. Balestriero says batch normalization actively concentrates regions around training samples, opposing the desired migration. Weight decay favors parameters near zero and generally promotes L2-style smoothness, while the robustness endpoint is closer to piecewise constant.

  • Increasing regularization can therefore slow adversarial grokking, postpone it beyond a practical training horizon, or prevent it entirely. The result complicates the old prescription that deliberately making networks simpler must improve generalization: the relevant question is which geometry the regularizer selects.

  • For a fixed compute allowance, Balestriero would not automatically fit the largest possible model and spend the remainder on training. If robustness is the objective, “you may want to use a much smaller model but allocate FLOPs for training time instead”—a direct argument for property-specific, “smart” FLOP allocation.

  • The geometry also suggests accelerators. Distance from a point to the nearest regional boundary is fast to compute and differentiable, so it can become a training regularizer. Architectural constraints offer another route: removing biases forces central, cone-shaped partitions. Dataset size, label noise, curricula, and teacher-student training may likewise shift grokking’s arrival.

7. Local complexity reveals training progress before accuracy does

  • The proposed local-complexity proxy counts how many partition regions lie near a point. Exact enumeration is infeasible for large networks, so the method counts per-layer vertices inside an epsilon ball and validates that statistic as a proxy for nearby regions.

  • Crucially, the measure peaks when clean train and test accuracy plateau, then begins declining well before adversarial grokking happens. It is therefore more sensitive to continued geometric change than headline metrics: the model can look finished while its internal partition is still reorganizing.

  • Balestriero proposes using the signal for early stopping, hyperparameter changes, or architecture selection—even without computing train or test accuracy. He cites prior work showing that partition statistics alone can support neural architecture search, suggesting geometry contains enough information to compare candidate networks.

8. Reconstruction learns visible energy before useful semantics

  • The reconstruction paper starts from two empirical mismatches. Autoencoder representations are useful baselines but often need downstream fine-tuning, and visually convincing reconstructions can appear long before the learned features become effective for perception.

  • Balestriero’s explanation runs through the image eigenspectrum. With pixel-space mean-squared error, high-energy components produce the largest gradients and the greatest immediate loss reduction. Those components are predominantly low frequency, so gradient descent learns blurry global structure before lower-amplitude, high-frequency detail.

  • The visual specimen carries the argument: low-frequency images remain blurry and hard to classify, while isolated high-frequency structure can make the class recognizable. The latter arrives late only because it contributes less to mean-squared error—not because it is less valuable to recognition.

  • High-frequency features also contain fewer easy background shortcuts: the network must attend to object shape rather than infer a label from grass or a beach. This is task- and dataset-dependent. On MNIST or SVHN, reconstruction and recognition align relatively well; with ImageNet-scale color, backgrounds, and object variation, the mismatch becomes severe.

9. Designed corruption can repair reconstruction’s spectral bias

  • Denoising autoencoders and masked autoencoders change the gradient diet by corrupting an image and asking the model to recover the original. Isotropic Gaussian noise and large masked blocks are different strategies, and a carefully chosen noise strategy can make selected parts of the inherited dataset bias harder to copy.

  • If practitioners know which frequencies matter, they can reverse-engineer a noise spectrum—the host mentions pink noise as an example—to suppress useless features and emphasize relevant ones. Balestriero’s caveat is that this requires expert knowledge about the downstream task.

  • The right corruption may not reduce to “high versus low frequency.” Depth estimation or counting trees could need different invariances, and an elaborate noise generator may erase the computational simplicity that made reconstruction attractive. Automatically discovering a useful, tractable corruption distribution remains open.

  • Reconstruction-free contrastive or non-contrastive methods compare augmented views in embedding space instead. Because they need not reproduce every pixel, they can discard nuisance detail while mapping different views to the same representation. Balestriero presents this as a more controllable proxy, not proof that all reconstruction objectives are unusable.

10. A few spline features turn LLM geometry into a toxicity detector

  • Each standard LLM layer contains multi-head attention followed by an MLP block; the MLP is again a spline. The paper extracts seven simple geometric features per block describing the region occupied by a prompt. Even a 70B model yields only about 500 features rather than millions of raw activation dimensions.

  • Without supervision, a two-dimensional visualization already clusters these features by prompt modality—mathematics, law, medicine—and separates toxic from non-toxic prompts. The region geometry is therefore informative about semantic content even though the features were not hand-designed for those categories.

  • A linear head converts them into a low-latency detector. The host’s OmniToxic comparison puts a popular model, reportedly downloaded 1.2 million times in the preceding month, near 73.5% area under the curve; the spline probe on Llama 2 7B reaches 99.18%. Balestriero confirms the comparison while emphasizing the method’s flexibility.

  • Using only the first three layers trades information for lower latency; one or two layers can go faster, while more layers can improve accuracy. Because the features are differentiable and computed on the fly, they could also drive data filtering, model comparison, regularization, new training objectives, or adversarial prompt manipulation.

11. High-dimensional prompts expose an RLHF coverage problem

  • For the attention block, the team derives a scalar proxy for the intrinsic dimension of a prompt’s subspace from attention sparsity. Longer contexts and more interrelated tokens reduce sparsity, increase effective dimension, and move the prompt into a region less represented during alignment training.

  • Artificially increasing that dimension produced toxic answers from models that rejected the original prompt. Balestriero’s explanation is extrapolative: RLHF teaches “don’t say that here” at sampled locations, but cannot automatically control every remote region in a gigantic high-dimensional space.

  • The discussion cites Llama 2 and an additional ChatGPT experiment, so Balestriero does not treat the result as architecture-specific. Padding with related natural-language concepts can work better than random tokens because it creates denser attention—and is much harder to detect than simply repeating a phrase many times.

  • His broader agenda at Brown is to replace empirical patching with accessible, industry-scale guarantees covering training dynamics, sparsity, regularization, and dataset bias. A failed method should yield a precise explanation, not “try another hyperparameter and come back to me in two days”; safe extrapolation will require better parameterizations or geometric control, not exhaustive sampling.