Neural Networks Learn Statistics of Increasing Complexity

ArXiv: 2402.04362

🎯 Pitch

Early in training, classifiers are so focused on simple statistical patterns that surgically transplanting only the means and covariances of one image class onto another tricks them into confidently assigning the wrong label—but this 'statistical myopia' completely vanishes as higher-order features are gradually incorporated, creating a striking U-shaped accuracy curve that peaks and then plummets as training progresses.

---

1. Executive Summary

This paper presents new empirical evidence for the distributional simplicity bias (DSB) by training neural networks on real image and text datasets, then probing their behavior throughout training on synthetic data that isolates sensitivity to low-order versus high-order statistics. The experimental framework uses two operationalized criteria—optimal transport to graft the means and covariances of one class onto another (e.g., making a CIFAR-10 airplane statistically resemble a truck at the first two moments), and maximum-entropy sampling to delete higher-order information while preserving specific statistics (e.g., evaluating ConvNeXt V2 on Gaussian samples matching only class-conditional means and covariances). Across architectures including ConvNeXt V2, Swin Transformer V2, and RegNet-Y on CIFAR-10/CIFARNet/SVHN/Fashion-MNIST/MNIST, the authors find a U-shaped accuracy curve: models classify first-order and second-order synthetic images correctly early in training but lose this ability later as they become sensitive to higher-order statistics—and for Pythia language models evaluated on n-gram-sampled sequences, they find a "double descent" where loss initially mirrors the U-shape then declines again through in-context learning. The DSB holds robustly across model scales and architectures, establishing that networks learn statistics of monotonically increasing order but that this progression saturates early for simple datasets where low-order moments already produce realistic images.

2. Context and Motivation

The Core Problem: We Don't Know What Neural Networks Learn First

The fundamental puzzle this paper tackles is deceptively simple: when a neural network learns from data, what does it learn first, and what does it learn later? This matters because neural networks exhibit a deeply counterintuitive generalization behavior—they can perfectly memorize randomly labeled training data (Zhang et al., 2021) yet still generalize well to unseen examples when trained on real data, and it is even possible to "poison" networks so they achieve zero training loss while behaving randomly on held-out test data (Huang et al., 2020). Understanding what networks latch onto during different phases of training is essential for explaining why they generalize at all, and why they sometimes fail catastrophically under distribution shift.

The distributional simplicity bias (DSB) offers a candidate explanation: neural networks learn to exploit the lower-order statistics of the input distribution first—the means and covariances—before progressively incorporating higher-order statistics like skewness and kurtosis. If true, this would mean that early-training networks are essentially "moment-limited" classifiers that see the world through a blurry statistical lens, gradually sharpening their perception as training proceeds. This has profound implications: it would explain why lightweight data augmentation helps early in training, why early-stopped models often transfer better, and why models can appear to understand a concept while being blind to its fine-grained structure.

However, prior to this work, the DSB was supported primarily by experiments on synthetic, theoretically tractable datasets. The paper's core gap is that no one had systematically tested whether real neural networks trained on real datasets actually behave as if they learn statistics in monotonic order. The distinction matters: synthetic datasets can be engineered to isolate statistical orders cleanly, but real datasets like CIFAR-10 or language corpora have entangled statistics where low-order moments may already capture substantial semantic content. Showing DSB on synthetic data leaves open the question of whether the bias is strong enough to shape learning dynamics on the messy distributions that practitioners actually care about.

Why This Problem Matters

Theoretical significance. The DSB sits at the intersection of several major research threads in deep learning theory. It connects to the neural tangent kernel (NTK) literature (Jacot et al., 2018), which shows that infinitely wide networks learn functions whose complexity is determined by the eigenfunctions of the initialization kernel—eigenfunctions that often exhibit strong simplicity biases toward low-frequency or low-order features (Canatar et al., 2021; Bietti & Mairal, 2019). It also connects to the Frequency Principle (Xu et al., 2019b; Rahaman et al., 2019), the empirical observation that networks fit low-frequency Fourier components of a target function before high-frequency ones. The DSB can be seen as a distributional analog of the Frequency Principle: instead of decomposing the target function into frequency components, it decomposes the input distribution into its moments. Unifying these perspectives would provide a more complete picture of how neural network inductive biases operate.

Practical implications for deployment. If networks genuinely rely on low-order statistics early in training, this has direct consequences for when to stop training. Early-stopped models should be more robust to distribution shifts that preserve means and covariances but alter higher-order structure—and conversely, more vulnerable to shifts that alter low-order statistics. This matters for safety-critical applications (medical imaging, autonomous driving) where understanding the failure modes of models under distribution shift is essential.

Self-supervised learning and transfer. The DSB suggests a natural curriculum: models trained to match low-order statistics first, then progressively higher-order ones, might learn more robust representations. Understanding the natural learning order could inform the design of self-supervised objectives that align with or deliberately counteract this bias.

Prior Approaches and Their Shortcomings

The paper identifies and addresses specific limitations in several prior research threads:

Synthetic-only DSB evidence. The most direct predecessor is Refinetti et al. (2023), who provided evidence for the DSB by training networks on a sequence of synthetic datasets that act as increasingly precise approximations to real data—for instance, datasets where only the mean matches the real data, then mean and covariance, then higher-order moments. They showed that early checkpoints trained on these approximate datasets perform about as well on real data as checkpoints trained directly on real data. This is clever, but it has a critical limitation: the networks were trained on synthetic approximations, not on real data. This means the experimental setup doesn't answer whether the DSB is strong enough to shape learning dynamics when a network is trained conventionally on a real dataset. The synthetic data might create an artificial bottleneck that forces the network to rely on low-order statistics, rather than revealing a natural tendency.

This paper inverts Refinetti et al.'s experimental design entirely. Instead of training on synthetic data and testing on real data, the authors train on real datasets (CIFAR-10, MNIST, the Pile) and test on synthetic data that probes sensitivity to statistics of different orders. This is a more direct assay of what real training dynamics produce. As the authors state:

"We believe this experimental design provides more direct evidence about the generalization behavior of commonly used models and training practices."

The Frequency Principle is related but distinct. Xu et al. (2019b) proposed the Frequency Principle—the tendency of neural networks to first fit low-frequency Fourier components of a target function before high-frequency ones—and subsequent work has explored this extensively (Rahaman et al., 2019; Xu et al., 2019a; Basri et al., 2020; Xu & Zhou, 2021). The Frequency Principle focuses on the function being approximated, decomposing it spectrally. The DSB focuses on the data distribution, decomposing it into moments. While these are related (low-order moments often correspond to low-frequency structure), they are not equivalent. The paper explicitly positions the DSB as complementary:

"Our work is similar in that we also aim to connect neural network learning order to simple mathematical properties of the training data, though we use distributional statistics, rather than frequency."

This distinction matters because some datasets have clear low-frequency structure but complex high-order moments, or vice versa. Understanding which decomposition better predicts learning dynamics is an open empirical question.

Architectural simplicity bias explanations are incomplete. An influential line of work attributes neural network generalization not to optimization dynamics but to the parameter-function map: the space of functions that a given architecture can represent, under random initialization, is biased toward simple functions regardless of how those functions are found (Valle-Perez et al., 2018; Chiang et al., 2023). In this view, gradient descent is not doing anything special—it is merely finding functions that the architecture was already biased toward. Chiang et al. (2023) provided striking evidence for this by showing that non-gradient-based optimizers, including unbiased random search over parameterizations, still generalize from training to testing data.

The DSB is not incompatible with architectural simplicity bias, but it makes a stronger claim: the learning trajectory matters, not just the final function. If networks learn statistics in monotonic order, then early checkpoints should be qualitatively different from late checkpoints in their sensitivity to statistical interventions, even if both generalize. The architectural simplicity bias alone doesn't predict this temporal ordering. The paper's U-shaped accuracy curves (e.g., Figure 3) are direct evidence that the learning trajectory has structure that architectural bias alone doesn't capture.

Prior language modeling order work is suggestive but limited. Choshen et al. (2022) empirically studied the learning dynamics of neural language models by tracking which grammatical patterns networks learn across training, and comparing network behavior to n-gram language models. They found that neural LMs initially match unigram and bigram model behavior early in training, then diverge. This is consistent with the DSB but was not framed in terms of statistical moments, and the connection between n-gram statistics and embedding-space moments was not established theoretically. The paper's Theorem 2.1—proving an exact equivalence between token n-gram frequencies and the moments of one-hot encoded sequences—provides the missing theoretical link, allowing the DSB framework to be applied to discrete domains with the same rigor as continuous ones.

How This Paper Positions Itself

The paper makes four explicit theoretical contributions that together define its position relative to prior work:

1. A Taylor expansion motivation for the DSB (Section 2, Equation 2). The expected loss of a network on a data distribution can be expressed as a sum over all central moments of the distribution, weighted by the corresponding partial derivatives of the loss evaluated at the mean input. At initialization, the Taylor coefficients $\frac{1}{\alpha!}$ cause higher-order moment contributions to decay factorially fast. As training progresses, the derivatives become correlated with the moments, potentially inflating higher-order terms. The paper argues that this inflation should occur in roughly monotonic order—second-order terms become important first, then third-order, and so on—because earlier terms account for factorially more of the loss at initialization and are thus higher-priority targets for gradient descent. This is not presented as a rigorous proof, but rather as a principled motivation that connects the DSB to the analytic structure of the loss function.

2. Two operationalized criteria for sensitivity to order-$k$ statistics (Section 2.2). Criterion 1: grafting the first $k$ statistics of class B onto class A should cause the model to treat A as B. Criterion 2: deleting higher-order statistics (by replacing them with maximum-entropy surrogates) should not harm performance. These criteria are intuitive but not trivial to implement—the paper develops substantial methodological machinery (optimal transport for grafting, maximum-entropy sampling with hypercube constraints for deletion) to test them rigorously.

3. Efficient methods for producing statistical probes (Sections 2.3–2.4). The paper describes Gaussian optimal transport (surgically matching means and covariances between class-conditional distributions), coordinatewise quantile normalization (matching first-order marginals), and a novel gradient-based method for approximate maximum-entropy sampling under hypercube constraints. These methods are both computationally tractable and theoretically grounded, making the framework practical for researchers to adopt.

4. An extension of the DSB to discrete domains (Section 2.5, Theorems 2.1–2.2). By proving that token n-gram frequencies are exactly equivalent to the moments of one-hot encoded sequences, and that this equivalence is preserved under linear embedding maps, the paper bridges the gap between the continuous statistics formalism and discrete language modeling. This is not just a theoretical curiosity—it enables the same maximum-entropy probing methodology to be applied to language models, using $k$-gram autoregressive models as the maximum-entropy distributions matching order-$k$ statistics.

The paper thus positions itself not as competing with the Frequency Principle, NTK theory, or architectural simplicity bias explanations, but as providing a complementary decomposition—one that is particularly natural for understanding distribution-level sensitivity and that extends cleanly across modalities (images and text). The experimental inversion of Refinetti et al.'s setup and the theoretical connection to n-gram models for language are the key positioning moves that distinguish this work from prior DSB research.

3. Technical Approach

3.1 Reader Orientation

The "system" in this paper is an experimental framework that generates synthetic probing data—images or text sequences whose statistical properties are surgically controlled—and evaluates a neural network on that data at regular intervals throughout training. It solves the problem of measuring which statistical orders a network is sensitive to at each point in training by enforcing two operational criteria: the model should classify modified data according to its low-order statistics when those statistics are grafted from another class, and it should maintain high performance when higher-order statistics are replaced with maximum-entropy noise that preserves only the low-order structure. The solution is a set of optimal transport and maximum-entropy sampling methods that take real training data, extract specified moments (means, covariances, coskewness), and produce synthetic datasets that isolate those moments for evaluation.

3.2 Big-Picture Architecture (Diagram in Words)

The framework has four major components connected in an evaluation pipeline:

1. Real Dataset — standard image classification datasets (CIFAR-10, MNIST, etc.) or a text corpus (the Pile) providing the raw training signal. Networks are trained normally on this data; the synthetic probes never appear during training.
2. Statistics Extraction — for each class or token distribution, compute the desired low-order statistics (means, covariance matrices, coskewness tensors, or n-gram frequency tables) from the real data. These are the "target statistics" that the synthetic data must match.
3. Synthetic Data Generators — two families of methods. Optimal transport maps take real images and transform them so that their low-order statistics match those of a different class, producing edited images that preserve the high-order structure of the source image but have the target's low-order properties. Maximum-entropy samplers generate entirely synthetic data—Gaussian samples, hypercube-constrained samples, or n-gram autoregressive sequences—that match the target statistics but otherwise maximize entropy (i.e., contain no information beyond those statistics).
4. Checkpoint Evaluation Loop — the network being trained is periodically evaluated (at log-spaced training steps) on each synthetic dataset. The evaluation metric is classification accuracy (for optimal transport probes, using the target class as the ground-truth label) or cross-entropy loss (for maximum-entropy probes). The resulting curves across training steps reveal when and how strongly the network exploits each statistical order.

Information flows in one direction: real data → statistics extraction → synthetic data generation → network evaluation → learning curves. Crucially, the network never sees the synthetic data during training—this is purely a probing framework, not a data augmentation technique.

3.3 Roadmap for the Deep Dive

• First, the Taylor expansion motivation (Equation 2), because it provides the theoretical intuition for why we expect monotonic learning of moments and formally defines what "moments" mean in the context of network loss. This is the conceptual foundation that the rest of the paper operationalizes.
• Second, the two criteria (Section 2.2), because they define precisely what it means for a network to be sensitive to order-$k$ statistics. Every subsequent method is a way to test one of these criteria: optimal transport for Criterion 1, maximum-entropy sampling for Criterion 2.
• Third, the optimal transport methods (Section 2.3), which generate images whose first- and second-order statistics match a target class while preserving higher-order structure. This is the most theoretically sophisticated component and requires understanding Gaussian OT, coordinatewise quantile normalization, and bounded shift.
• Fourth, the maximum-entropy sampling methods (Section 2.4), which generate synthetic data from scratch matching only the specified statistics. This includes Gaussian sampling, Conrad distributions, hypercube-constrained optimization, and independent coordinate sampling. The gradient-based approximate sampling method is a novel technical contribution.
• Fifth, the discrete-domain extension (Section 2.5, Theorems 2.1–2.2), which proves that n-gram statistics are equivalent to embedding-space moments and thereby enables the same probing methodology for language models. This is what makes the paper cross-modal rather than vision-only.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a methodological and empirical paper whose core idea is that the distributional simplicity bias can be operationalized through two criteria—grafting low-order statistics and deleting higher-order ones—and tested using optimal transport and maximum-entropy sampling across both continuous (image) and discrete (language) domains.

---

The Taylor Expansion Motivation for the DSB

The paper does not present a rigorous proof that networks must learn moments in monotonic order. Instead, it provides a Taylor series argument that makes monotonic learning a natural consequence of the structure of the expected loss, given mild assumptions about how derivatives evolve during training. This argument is the intellectual scaffolding for the entire experimental framework.

The starting point is the loss function $L(x)$ of a neural network evaluated on a single input $x \in \mathbb{R}^d$. If $L$ is analytic (infinitely differentiable with a convergent Taylor series) around the mean input $\mu = \mathbb{E}[x]$, we can expand it as:

$L(x) = \sum_{\alpha \in \mathbb{N}^d} \frac{(x - \mu)^\alpha}{\alpha!} (\partial^\alpha L)(\mu)$

where $\alpha = (\alpha_1, \alpha_2, \ldots, \alpha_d)$ is a multi-index—a tuple of $d$ non-negative integers, one per input coordinate. The expression $(x - \mu)^\alpha$ denotes the product $(x_1 - \mu_1)^{\alpha_1} (x_2 - \mu_2)^{\alpha_2} \cdots (x_d - \mu_d)^{\alpha_d}$, $\alpha!$ denotes the product of factorials $\alpha_1! \alpha_2! \cdots \alpha_d!$, and $(\partial^\alpha L)(\mu)$ is the mixed partial derivative $\frac{\partial^{|\alpha|} L}{\partial x_1^{\alpha_1} \partial x_2^{\alpha_2} \cdots \partial x_d^{\alpha_d}}$ evaluated at $\mu$, where $|\alpha| = \sum_{i=1}^d \alpha_i$ is the order of the derivative.

What this equation computes: the loss at any input $x$ as an infinite weighted sum over all monomials $(x - \mu)^\alpha$, where the weight for each monomial is the corresponding partial derivative of the loss at the mean, divided by the factorial of the multi-index. The zeroth-order term ($|\alpha| = 0$) is simply $L(\mu)$, the loss at the mean. The first-order terms ($|\alpha| = 1$) are $\frac{\partial L}{\partial x_i}(\mu) \cdot (x_i - \mu_i)$ for each coordinate $i$. The second-order terms ($|\alpha| = 2$) include both pure quadratic terms $\frac{1}{2}\frac{\partial^2 L}{\partial x_i^2}(\mu) \cdot (x_i - \mu_i)^2$ and mixed quadratic terms $\frac{\partial^2 L}{\partial x_i \partial x_j}(\mu) \cdot (x_i - \mu_i)(x_j - \mu_j)$ for $i \neq j$.

Why this form: the Taylor expansion decomposes the loss into contributions from different "orders" of deviation from the mean. The key property is that the factorial $\alpha!$ in the denominator grows super-exponentially with the order $|\alpha|$. For a multi-index of order $k$, the factorial is at least $k!$ (when all components are zero except one, which equals $k$), and can be much larger when the order is distributed across multiple coordinates. This factorial decay means that, all else being equal, higher-order monomials contribute less to the sum than lower-order ones.

Now take the expectation of both sides over the data distribution. Because expectation is linear (it distributes over sums), we get:

$\mathbb{E}[L(x)] = \sum_{\alpha \in \mathbb{N}^d} \frac{(\partial^\alpha L)(\mu)}{\alpha!} \underbrace{\mathbb{E}[(x - \mu)^\alpha]}_{\text{central moment}}$

What this equation computes: the expected loss of the network on a data distribution, expressed as an infinite sum over all central moments of that distribution, each multiplied by the corresponding normalized derivative of the loss at the mean. The central moment $\mathbb{E}[(x - \mu)^\alpha]$ is a scalar that characterizes the distribution's shape: $|\alpha| = 1$ gives the first central moments (which are exactly zero by definition of $\mu$), $|\alpha| = 2$ gives the covariance structure (variances and cross-covariances), $|\alpha| = 3$ gives coskewness, $|\alpha| = 4$ gives cokurtosis, and so on.

Why this form: this equation directly connects the data distribution (through its moments) to the network's loss (through its derivatives). The contribution of each statistical order $k$ to the expected loss is the sum over all multi-indices with $|\alpha| = k$ of $\frac{(\partial^\alpha L)(\mu)}{\alpha!} \mathbb{E}[(x - \mu)^\alpha]$. If the derivatives $(\partial^\alpha L)(\mu)$ are all roughly the same magnitude at initialization—which is plausible for randomly initialized networks—then the $1/\alpha!$ factor causes higher-order moment contributions to be factorially suppressed relative to lower-order ones. The expected loss is initially dominated by low-order statistics.

The DSB intuition then follows from considering how training changes these terms. During training, gradient descent adjusts the network parameters to reduce $\mathbb{E}[L(x)]$. The paper argues that:

1. The factorial coefficients make low-order terms larger in magnitude at initialization, so they account for more of the initial loss. Gradient descent, which greedily reduces loss, will therefore prioritize reducing these larger terms—meaning it will learn to exploit low-order statistics first.

2. As training proceeds, the derivatives $(\partial^\alpha L)(\mu)$ become correlated with the corresponding central moments. This correlation can inflate the magnitude of higher-order terms beyond what the $1/\alpha!$ factor would suggest. The paper argues this inflation should occur in roughly monotonic order—second-order terms become important first, then third-order, and so on—because earlier terms were both larger at initialization and are thus "higher-priority targets" for the optimizer until they can no longer be easily reduced, at which point the optimizer moves on to the next order.

A critical auxiliary assumption is that inputs are constrained to the unit hypercube $[0, 1]^d$. Since every coordinate of $x$ is between 0 and 1, the product $(x_i - \mu_i)^{\alpha_i}$ has magnitude at most $1^{\alpha_i} = 1$, and the moment $\mathbb{E}[(x - \mu)^\alpha]$ is also bounded by 1. In fact, when coordinates are independent, higher-order moments decay exponentially with order, further reinforcing the initial dominance of low-order terms. Standard PyTorch preprocessing maps image pixel intensities to $[0, 1]$, and one-hot encoded token sequences satisfy this constraint by construction, so the assumption holds for both modalities studied in the paper.

The paper is careful to note that ReLU networks are not analytic (the ReLU function is not differentiable at zero), so the Taylor expansion is not strictly valid. However, it is possible to construct arbitrarily close analytic approximations to ReLU (e.g., GELU, as in Hendrycks & Gimpel, 2016), and the paper treats the Taylor expansion as a motivating intuition rather than a formal proof. The experimental results are what ultimately validate (or refute) the DSB.

---

The Two Operational Criteria

The Taylor expansion tells us that the expected loss depends on moments of all orders, but it does not tell us which moments the network is actually using to make predictions. A network could theoretically have large second-order derivative terms but still rely primarily on first-order statistics for classification—the derivatives reflect sensitivity of the loss, not necessarily the decision boundary.

To make the DSB testable, the paper needs criteria that translate "the network uses order-$k$ statistics" into experimentally measurable quantities. It proposes two:

Criterion 1 (Grafting): If the low-order statistics (up to order $k$) of images from class A are surgically replaced with those of class B, while preserving the higher-order structure of class A, a network sensitive to statistics up to order $k$ should classify the resulting images as class B. The network's decision boundary has been "fooled" into following the grafted statistics.

The operational question this criterion answers: does changing the first $k$ moments change the classification? If the answer is yes, the network must be using those moments (or at least be sensitive to them). If no, the network is either insensitive to those moments (relying on something else) or is using higher-order structure that overrides the grafted statistics.

Criterion 2 (Deletion): If all information beyond order $k$ is deleted from the data—replaced with a maximum-entropy distribution that matches the first $k$ moments exactly but contains no other structure—a network that relies only on statistics up to order $k$ should achieve the same performance as on real data. If performance drops, the network was exploiting higher-order statistics beyond $k$.

The operational question: is performance preserved when higher-order information is removed? The principle of maximum entropy is crucial here because it provides a principled way to "delete" information: among all distributions matching the specified moments, the maximum-entropy distribution is the least informative one—it introduces no correlations or structure beyond what is forced by the moment constraints. If the network's performance on maximum-entropy samples equals its performance on real data, then the real data contained no useful information beyond those moments for this network at this point in training.

These two criteria are complementary. Criterion 1 tests whether the network uses the low-order statistics—it's an active intervention that should produce a specific behavior change. Criterion 2 tests whether the network needs anything beyond the low-order statistics—it's a deletion test that should produce no behavior change if the network is moment-limited. Together, they provide converging evidence: if grafting changes classification and deletion doesn't harm performance, the network is genuinely operating at that statistical order.

The rest of the technical approach consists of building the machinery to test these criteria. Optimal transport provides the "grafting" operation for Criterion 1. Maximum-entropy sampling provides the "deletion" operation for Criterion 2. The theoretical extension to discrete domains shows that the same criteria apply to language models, where the "moments" are n-gram frequencies.

---

Optimal Transport Methods for Statistical Grafting (Criterion 1)

Optimal transport theory provides a principled framework for moving probability mass from one distribution to another while minimizing a cost function—typically the expected squared Euclidean distance that samples are moved. Given a source distribution $P$ and a target distribution $Q$, an optimal transport map $T: \text{supp}(P) \to \text{supp}(Q)$ is a function that pushes $P$ forward to $Q$ (meaning $T_\sharp P = Q$, where $T_\sharp$ denotes the pushforward: the distribution of $T(x)$ when $x \sim P$ is exactly $Q$) while minimizing $\mathbb{E}_{x \sim P}[\|x - T(x)\|_2^2]$.

The paper uses three specific optimal transport methods, each affecting different orders of statistics:

Gaussian Optimal Transport

This is the workhorse method for testing Criterion 1 with second-order statistics (means and covariances). Given two distributions $P$ and $Q$, the map:

$T(x) = A(x - \mu_P) + \mu_Q$

where $A$ is a matrix defined as:

$A = \Sigma_P^{-1/2} \left(\Sigma_P^{1/2} \Sigma_Q \Sigma_P^{1/2}\right)^{1/2} \Sigma_P^{-1/2}$

is the optimal transport map from $P$ to $Q$ when $P = \mathcal{N}(\mu_P, \Sigma_P)$ and $Q = \mathcal{N}(\mu_Q, \Sigma_Q)$ are both Gaussian. More generally, even when $P$ is not Gaussian but has finite second moments, applying $T$ to samples from $P$ transports them to a distribution with mean exactly $\mu_Q$ and covariance exactly $\Sigma_Q$, and this map minimizes the expected squared $L_2$ transport cost (Dowson & Landau, 1982).

What it computes: for each image $x$ drawn from the source class, subtract the source class mean $\mu_P$, apply the linear transformation $A$ that "whitens" the source covariance and then "recolors" it with the target covariance, and add the target class mean $\mu_Q$. The result is an image that looks structurally like the original (because the transformation is linear and minimizes pixel displacement) but whose first two moments now match the target class.

Implementation details: Given $k$ image classes, each containing tensors of shape $C \times H \times W$ (channels × height × width), the images are first unrolled into vectors of size $d = CHW$. For each class, the empirical mean $\mu \in \mathbb{R}^d$ and covariance matrix $\Sigma \in \mathbb{R}^{d \times d}$ are computed from the training images. Because the covariance matrix has $d^2$ elements and $d$ can be large (e.g., for 32×32×3 CIFAR-10 images, $d = 3072$, so $\Sigma$ has about $9.4 \times 10^6$ entries), the paper applies linear shrinkage (Bodnar et al., 2014) to improve the covariance estimate—this regularizes the sample covariance toward a scaled identity matrix, which improves numerical stability and estimation accuracy when the number of samples is not vastly larger than $d$.

The matrix $A$ involves matrix square roots and inverses of $\Sigma_P$ and $\Sigma_Q$, which are computed via eigendecomposition. For each pair of classes $(i, j)$, the map $T_{i \to j}$ is computed from class $i$ to class $j$. At evaluation time, each test image from class $i$ is passed through $T_{i \to j}$, and the network's prediction is compared against the label for class $j$ (the target class, not the source class).

Why this form: The appeal of Gaussian OT is that it provides a closed-form, differentiable map that exactly matches means and covariances while minimally perturbing the images. Because the map is linear, it preserves the "shape" of the source distribution's deviations from its mean—if the source class has distinctive high-order texture patterns, those patterns survive the transformation, only now centered at the target mean and stretched/rotated to match the target covariance. This means that if a network classifies the transformed image as the target class, it must be relying primarily on the mean and covariance, because the higher-order structure is still from the source class.

Computational cost: The method requires $O(d^3)$ compute and $O(d^2)$ memory due to the eigendecomposition of the covariance matrix, where $d$ is the number of pixels. Since $d$ grows as $H \times W \times C$, this scales with the fourth power of image resolution—so the method is limited to $32 \times 32$ or $64 \times 64$ images. The CIFARNet dataset was specifically constructed at $64 \times 64$ to stay within this computational envelope while providing higher-resolution images than CIFAR-10. The paper notes that NumPy and PyTorch eigensolvers would crash on very large covariance matrices (PyTorch issue #92141), imposing a further practical constraint.

Coordinatewise Quantile Normalization (CQN)

This is a first-order method that matches only the marginal (per-pixel) distributions between classes, approximately grafting the means while leaving correlations largely intact. It works by applying the probability integral transform coordinate by coordinate.

Algorithm: For a given image coordinate (a specific pixel across all images in a class), let $F_X$ be its empirical cumulative distribution function (CDF) in the source class, and $F_Y$ be its empirical CDF in the target class. For each pixel intensity $x$ in a source image, the transformation $F_Y^{-1}(F_X(x))$ is applied, where $F_Y^{-1}$ is the empirical quantile function (inverse CDF) of the target class at that coordinate.

What it computes: First, $F_X(x)$ maps the pixel intensity to its percentile rank within the source class—what fraction of source images have a lower intensity at this pixel? This yields a value in $[0, 1]$. Then $F_Y^{-1}$ maps that percentile back to a pixel intensity—what intensity in the target class corresponds to that same percentile? The result is that the marginal distribution of each pixel coordinate now matches the target class exactly, because the transformation is a monotone increasing function that pushes the source distribution to the target distribution coordinatewise.

Why this form: Quantile normalization is known to be the optimal transport map for a large class of cost functions when the cost depends only on the values of individual coordinates (Santambrogio, 2015, Ch. 2.2). It is "first-order" in the sense that it matches all single-coordinate moments (the homogeneous moments $\mathbb{E}[x_i^k]$ for each $i$ and all $k$) but does not enforce joint moments like covariances $\mathbb{E}[x_i x_j]$ for $i \neq j$. The fact that the Pekinese dog in Figure 1 (center) remains recognizable after CQN to match the goldfish class illustrates that correlations between pixels are largely preserved—the dog shape persists even though the color palette has shifted.

Practical advantage: CQN has negligible computational cost compared to Gaussian OT—it requires only computing and storing per-pixel empirical CDFs from the training data, then evaluating them per image. It serves as a baseline for what first-order information alone can achieve.

Bounded Shift

This is an alternative first-order method that surgically changes the mean of each image while respecting the $[0, 1]$ pixel intensity constraint. Unlike CQN, which matches the entire marginal distribution, bounded shift only changes the per-image mean to match a target value, making the minimal $L_2$ perturbation to the pixel intensities subject to the box constraint.

Algorithm (Appendix D, Algorithm 1): Given an image vector $x \in [0, 1]^n$ (where $n$ is the number of pixels after unrolling) with current mean $\bar{x} = \frac{1}{n}\sum_{i=1}^n x_i$, and a desired mean $m$:

1. Sort the coordinates of $x$ in descending order (largest pixel intensities first).
2. For each coordinate $i$ from largest to smallest:
   • Compute the candidate shift: $y_i = x_i + (m - \bar{x})$. This would add the same constant to every pixel, uniformly shifting the mean by $m - \bar{x}$.
   • If $y_i \leq 1$, this candidate is valid (doesn't exceed the maximum intensity). Keep it, and all subsequent smaller coordinates will also be valid because they start smaller.
   • If $y_i > 1$, the candidate overshoots. Clamp $y_i = 1$ (maximum allowable intensity), update the remaining target mean to $m' = \frac{nm - i}{n - i}$ (redistributing the "excess" mean that couldn't be added to this pixel across the remaining pixels), and continue to the next coordinate.
3. After processing all coordinates, restore the original ordering (unsort) to get the final edited image.

What it computes: the solution to the optimization problem $\min_{y \in [0,1]^n} \|x - y\|_2^2$ subject to the constraint that $\frac{1}{n}\sum_i y_i = m$. It adds a constant shift to all pixels, but whenever a pixel would exceed 1, it caps it and redistributes the remaining required mean increase to the smaller (darker) pixels, which still have room to increase.

Why this form: This is the exact optimal solution under $L_2$ cost with the box constraint, proven in Theorem D.1 of the paper. It is complementary to CQN: CQN matches the full marginal distribution (all moments of individual coordinates), while bounded shift matches only the mean but preserves relative intensity ordering within each image. The paper includes both to disentangle whether networks are sensitive to the mean specifically or to the full marginal distribution. The proof in Theorem D.1 uses KKT conditions and exploits the reflection symmetry of the problem (the case of decreasing the mean $m < \bar{x}$ reduces to the increasing case by reflecting all values about $1/2$).

---

Maximum-Entropy Sampling for Statistical Deletion (Criterion 2)

The principle of maximum entropy (Jaynes, 1957) states that among all probability distributions satisfying a set of known constraints, the one that maximizes entropy is the "least informative" distribution consistent with those constraints—it introduces no additional structure beyond what is forced by the constraints. For Criterion 2, the constraints are that the synthetic data must match the first $k$ moments (or other specified low-order statistics) of the real data distribution. The maximum-entropy distribution satisfying these constraints then represents what the data would look like if all higher-order structure were deleted.

The paper uses several specific maximum-entropy distributions at different constraint levels, each testing a different statistical order.

Gaussian Sampling (Second-Order, Unconstrained Support)

The maximum-entropy distribution on $\mathbb{R}^d$ with known mean $\mu \in \mathbb{R}^d$ and covariance matrix $\Sigma \in \mathbb{R}^{d \times d}$ (and no other constraints) is the multivariate Gaussian $\mathcal{N}(\mu, \Sigma)$. This is a standard result in information theory: the Gaussian maximizes entropy subject to first- and second-moment constraints.

What it produces: samples that match the class-conditional means and full covariance structure of the real images, but contain no higher-order structure—no skewness, no kurtosis, no higher-order correlations. All pixel values are real numbers, not constrained to $[0, 1]$.

Why this matters for the DSB test: If a network achieves high accuracy on Gaussian samples that match only the class-conditional means and covariances, it cannot be relying on skewness, kurtosis, or any other higher-order statistic of the real data distribution. The Gaussian samples are the "most random" distribution with those first two moments. Any excess performance over random guessing must be attributable solely to mean and covariance information.

Practical issue: Gaussian samples do not respect the $[0, 1]$ pixel intensity constraint of real images, and can produce values far outside this range (especially in high dimensions, where the typical set is a thin shell far from the mean in Euclidean distance, even though each coordinate is centered near its mean). The network might fail on these samples not because it needs higher-order statistics, but because the support violation creates an out-of-distribution input that the network hasn't learned to handle. This motivates the hypercube-constrained methods below.

The Conrad Distribution (First-Order, Hypercube-Constrained)

For a single coordinate with known mean $m \in [0, 1]$ and no other constraints, the paper derives the maximum-entropy distribution supported on $[0, 1]$. The density is:

$p(x) = \frac{b \exp(-bx + b)}{\exp(b) - 1}$

where $b \neq 0$ is a parameter chosen so that the mean equals $m$, specifically satisfying:

$m = \frac{-b - \exp(b) + 1}{b(\exp(b) - 1)}$

In the special case $m = 1/2$, the maximum-entropy distribution is $\text{Unif}(0, 1)$ (the uniform distribution), and $b \to 0$.

What it computes: for a given target mean $m$, Newton's method (or another root-finding algorithm) solves the mean equation for $b$. Once $b$ is known, inverse transform sampling generates samples: sample $U \sim \text{Unif}(0, 1)$, then solve $F(x) = U$ for $x$, where $F$ is the CDF of the Conrad distribution. The resulting samples are bounded in $[0, 1]$ and have exactly the desired mean.

Why this form: the derivation (Theorem E.1) follows the standard maximum-entropy recipe: the density must be of the form $\exp(-\lambda_0 - \lambda_1 x)$ (exponential family with the mean as sufficient statistic), where the Lagrange multipliers $\lambda_0, \lambda_1$ are chosen to satisfy normalization and mean constraints. Solving the constraints yields the parameter $b = \lambda_1$. The paper notes that this formula is "not well-known" despite an alternative derivation in Conrad (2004), so providing it here is a minor theoretical contribution in its own right.

Usage: To generate first-order maximum-entropy images, a Conrad distribution is fit to the mean of each pixel coordinate independently. Synthetic images are then generated by sampling each pixel from its fitted Conrad distribution, treating coordinates as independent. This matches the first-order statistics (per-pixel means) exactly, matches no higher-order mixed moments, and respects the $[0, 1]$ constraint. The paper calls this "Conrad sampling" in Figure 5.

Independent Coordinate Sampling (ICS)

ICS takes a slightly different approach to first-order matching. Instead of deriving a parametric maximum-entropy distribution, it directly matches the empirical marginal CDF of each coordinate, then samples each coordinate independently.

Algorithm: For each pixel coordinate across all training images of a class, compute the empirical CDF $\hat{F}_i$ (essentially the sorted list of observed intensities). To generate a synthetic image, for each coordinate $i$, sample $U \sim \text{Unif}(0, 1)$ and set the pixel value to $\hat{F}_i^{-1}(U)$—a randomly chosen intensity from the empirical distribution of that coordinate.

What it computes: samples where each pixel individually follows the real marginal distribution, but pixels are independent. This matches all homogeneous moments $\mathbb{E}[x_i^k]$ for each coordinate $i$ and all orders $k$ (since the full marginal is matched), but matches no mixed moments $\mathbb{E}[x_i x_j]$ for $i \neq j$. In high dimensions, almost all moments above first order are mixed rather than homogeneous (for $d = 3072$, there are $d$ first-order moments but $\binom{d}{2} \approx 4.7 \times 10^6$ second-order mixed moments), so ICS effectively matches first-order statistics and almost nothing beyond.

Why this form: ICS is conceptually simple and computationally trivial—it requires only storing empirical CDFs per pixel and sampling from them. It provides a complementary first-order baseline to Conrad sampling: Conrad maximizes entropy given only the mean (so pixel distributions may differ from the real marginals), while ICS exactly matches the real marginals (so all single-pixel moments are preserved). If networks behave similarly on both, it confirms that only the means matter; if they differ, the full marginal distributions carry useful information beyond the mean.

Hypercube-Constrained Maximum-Entropy Sampling (Second- and Third-Order)

The Gaussian matches means and covariances but violates the $[0, 1]$ support constraint. The maximum-entropy distribution with known mean and covariance and support $[0, 1]^d$ is known in the 1D case (it takes the form $p(x) = \exp(-\lambda_0 - \lambda_1 x - \lambda_2 x^2)$, which can be a truncated normal or a U-shaped distribution depending on the Lagrange multipliers; Dowson & Wragg, 1973), but there is no known analytic formula for the multidimensional case. The log-density would still be a quadratic form, but the normalization constant (partition function) is intractable, and the "scale matrix" may not be positive semidefinite.

The paper solves this with a novel gradient-based approximate sampling method. Rather than deriving the density analytically and sampling from it, the method directly optimizes a finite set of synthetic samples to have the desired statistics while maximizing an estimate of the population entropy.

Algorithm sketch:

1. Initialize a batch of synthetic images randomly (e.g., uniform noise in $[0, 1]^d$).
2. Compute the current batch statistics—mean, covariance, and optionally coskewness—and compare them to the target statistics from the real data via a loss function (e.g., mean squared error between current and target moments).
3. Compute the Kozachenko-Leonenko estimate of the entropy of the empirical distribution of the batch. This is a non-parametric entropy estimator based on nearest-neighbor distances (Kozachenko & Leonenko, 1987; Sablayrolles et al., 2018) that penalizes samples being too close together (encouraging spread) or too regular (encouraging randomness).
4. Take a gradient step to adjust the synthetic images, optimizing a combined objective: minimize moment matching loss plus maximize entropy estimate.
5. Clamp pixel values to $[0, 1]$ after each gradient step.
6. Repeat for many iterations (e.g., 10,000 steps per class for CIFARNet third-order matching).

What it produces: a finite set of synthetic images whose sample moments closely match the target moments, whose pixel values are all in $[0, 1]$, and which are as "random" (high-entropy) as possible subject to those constraints. The method works for second-order constraints (matching means and covariances only) and third-order constraints (matching means, covariances, and coskewness tensors).

Why this form: This method bypasses the need for an analytic density, which is intractable for multidimensional hypercube-constrained maximum-entropy distributions. It directly produces the samples that an evaluator needs, without ever representing a density function. The entropy maximization term is crucial: without it, the optimization could simply memorize the real training images (which trivially match all target statistics) rather than producing "generic" maximum-entropy samples. The entropy term forces the synthetic images to spread out and avoid reproducing specific real images.

Computational considerations: The coskewness tensor for CIFARNet (64×64×3 = 12,288 pixels) has dimensions $12288 \times 12288 \times 12288$, which would require nearly 8 terabytes in full precision. The paper avoids computing it explicitly: each optimization step computes the coskewness only along randomly sampled slices of length $l$, meaning only two $12288 \times 12288 \times l$ tensors are stored at any time (one for real data, one for synthetic). The method still requires 10,000 steps per class, totalling about 36 hours on an NVIDIA A40 GPU for CIFARNet. For lower-resolution datasets (CIFAR-10, 32×32), the cost is much lower—roughly 65 seconds and 29 GB of GPU memory on an NVIDIA L40 for 200K synthetic images.

The paper's Figure 2 visually demonstrates the effect of these maximum-entropy distributions. When the underlying real images are simple (MNIST digits—essentially white strokes on a black background), the second-order hypercube-constrained samples are still clearly recognizable as digits (blurry but identifiable). For more complex distributions (CIFAR-10), the second-order samples look like colored noise—the means and covariances capture the average color palette and blurry shape, but none of the detailed structure needed for human recognition. This suggests that CIFAR-10 requires higher-order statistics for good performance, while MNIST might be solvable with only first- and second-order information. The experimental results in Section 3.3 confirm this: MNIST and Fashion MNIST show much less U-shaped non-monotonicity in their accuracy curves than CIFAR-10.

---

Extension to Discrete Domains: The N-Gram / Moment Equivalence

The entire framework so far has been built for continuous inputs (images), where moments are standard mathematical objects. To apply the DSB to language modeling, the paper needs to connect the continuous notion of "moments" to the discrete notion of "n-gram statistics" that characterizes token sequence distributions. Theorems 2.1 and 2.2 provide this bridge.

One-Hot Encoding and Flattening

Consider a vocabulary $\mathcal{V}$ of size $|\mathcal{V}|$, and sequences of length $N$. A sequence of tokens is first converted to a sequence of one-hot vectors: each position $t$ is represented by a vector in $\{0, 1\}^{|\mathcal{V}|}$ with a 1 at the index of the token and 0 elsewhere. This $N \times |\mathcal{V}|$ Boolean matrix is then flattened (concatenated row by row) into a single vector in $\{0, 1\}^{N \cdot |\mathcal{V}|}$. The function $f: \mathcal{V}^N \to \{0, 1\}^{N \cdot |\mathcal{V}|}$ performs this encoding.

What this flattening does: it converts a discrete sequence into a point in a continuous (Boolean) hypercube, exactly analogous to an image flattened into a vector of pixel intensities. Now the standard machinery of moments applies: for any multi-index $\alpha \in \mathbb{N}^{N \cdot |\mathcal{V}|}$, the moment $\mathbb{E}[f(x)^\alpha]$ is defined.

Theorem 2.1: N-Gram Statistics ARE Moments

The theorem states: Every moment of the one-hot encoded distribution $f_\sharp P$ is equal to an n-gram statistic of $P$, and vice versa.

Proof sketch: Each component of the flattened vector corresponds to a specific (sequence position, vocabulary token) pair. Because the components are Boolean (0 or 1), a monomial $f(x)^\alpha$ is 1 if and only if, for every index $i$ where the multi-index $\alpha_i > 0$, the corresponding token is present at the corresponding position. This is exactly a logical conjunction: the token at position $t_1$ is $v_1$ AND the token at position $t_2$ is $v_2$ AND ... AND the token at position $t_k$ is $v_k$, where $k = |\alpha|$ (the number of non-zero entries in $\alpha$). The expectation of this Boolean product is simply the probability that all these conjunctions hold—which is precisely an n-gram statistic: the probability that a specific set of tokens appears at a specific set of positions.

The paper notes two non-standard aspects of its "n-gram statistic" definition:

1. It includes skip-grams (e.g., "the _ dog", where the underscore matches any token at that position), because the multi-index can have zeros at some positions.
2. It is tied to absolute sequence positions, not relative offsets. For the Pythia models used in experiments, this is reasonable because they were trained on fixed-length chunks of text sampled from longer documents (Biderman et al., 2023), so absolute position should not significantly affect n-gram probabilities. The paper assumes translation invariance in what follows.

A subtle point: there are infinitely many moments corresponding to the same n-gram, because any multi-index with entries greater than 1 is equivalent to the binarized version (since $x_i^2 = x_i$ for Boolean $x_i \in \{0, 1\}$). The paper acknowledges this redundancy and treats the binarized multi-index as the canonical representation.

Theorem 2.2: Equal Embedding Moments

Given an embedding matrix $E \in \mathbb{R}^{|\mathcal{V}| \times d}$ that maps one-hot token vectors to $d$-dimensional embedding vectors, two distributions over token sequences with identical n-gram statistics up to order $k$ will produce embedded sequences with identical moments up to order $k$.

Proof sketch: Embedding is a linear operation (multiplying the one-hot matrix by $E$). Composing the one-hot encoding $f$ with the embedding multiplication $g$ (which applies $E$ to each row of the one-hot matrix and flattens the result) gives the embedding-space representation. Since $g$ is linear, $g(x)^\alpha$ is a polynomial in the components of the one-hot vector $x$ with degree equal to $|\alpha|$. The expectation $\mathbb{E}[(g \circ f)(x)^\alpha]$ is therefore a linear combination of moments of $f_\sharp P$ of order at most $|\alpha|$. If two distributions have identical moments up to order $k$, their embedded distributions also have identical moments up to order $k$.

Why these theorems matter: They provide the theoretical justification for testing the DSB in language models using the same maximum-entropy methodology as images. The "moments" of the token distribution are n-gram frequencies. The maximum-entropy distribution matching n-gram frequencies up to order $k$ is exactly a $k$-gram autoregressive language model—a Markov model where each token's probability depends only on the preceding $k-1$ tokens (for standard n-grams; the paper's definition with absolute positions and skip-grams means the model conditions on any specified set of previous positions, but the standard bigram/trigram approximations used in experiments are standard).

Theorem 2.2 further guarantees that this equivalence holds in embedding space, which is where the neural network actually operates. Even though the network receives discrete tokens as input, its internal representations live in the continuous embedding space, and the moments in that space are what matter for the Taylor expansion argument. The theorem assures us that matching n-gram statistics forces matching embedding-space moments, so the maximum-entropy n-gram sequences are valid probes for the DSB in language models.

Experimental application: For Pythia language models, the paper computes token unigram and bigram frequencies across the training corpus (the Pile), constructs unigram and bigram autoregressive models that sample each token independently (unigram) or conditionally on the previous token (bigram), and generates synthetic sequences of length 2049 from these models. These sequences are then fed to Pythia checkpoints throughout training to measure cross-entropy loss. A $k$-gram model is the maximum-entropy distribution matching all n-gram statistics up to order $k$ (for standard contiguous n-grams without skip-grams—the paper acknowledges that skip-grams would make the maximum-entropy model more complex and omits them for tractability). The same logic extends to 3-gram and 4-gram models evaluated on a subset of the Pile.

4. Key Insights and Innovations

Innovation 1: Inverting the DSB Experimental Paradigm from Training to Probing

The single most distinctive intellectual move in this paper is its experimental inversion of Refinetti et al. (2023). Instead of asking "if we train a network on a low-order approximation of real data, does it behave like a network trained on real data?", the paper asks "if we train a network on real data and then probe it with low-order synthetic data, does it behave as if it only sees low-order statistics?" The difference sounds subtle but is fundamental: the former tests whether low-order statistics are sufficient to produce the behavior of a network trained on real data; the latter tests whether a network trained on real data actually relies on low-order statistics at a given point in training.

This inversion matters because Refinetti et al.'s design cannot disentangle whether the DSB is a property of the learning trajectory or an artifact of the training data bottleneck. If you train exclusively on a dataset that contains only means and covariances, the network has no choice but to rely on them—it would be surprising if it didn't. But when you train on the full real distribution (with all higher-order structure present) and never expose the network to synthetic data during training, the probe results reveal genuine learning dynamics. The network is free to exploit any statistics available at any training step; the fact that it nevertheless classifies Gaussian-transported images by their grafted means and covariances—and loses this ability later—means the DSB is a real temporal bias in how networks extract structure from data, not a constraint imposed by data starvation.

The paper's explicit framing of this design choice makes the contribution clear:

"We build on Refinetti et al. (2023) by inverting their experimental setup. We train our models on real datasets, then test them throughout training on synthetic data that probe the model's reliance on statistics of different orders. We believe this experimental design provides more direct evidence about the generalization behavior of commonly used models and training practices."

This is not a minor variant—it is a more stringent test that closes a loophole in prior DSB evidence. The probing methodology is what allows the paper to observe the characteristic U-shaped accuracy curves (Figure 3), where networks temporarily classify maximum-entropy and OT-edited images better than random, then regress. Training on synthetic data would not produce this non-monotonicity because the network would have no opportunity to "move beyond" low-order statistics. The U-shape is the paper's signature finding and it is only visible because of the inversion.

Innovation 2: Operationalizing "Statistical Reliance" via Optimal Transport and Maximum Entropy

Before this work, the DSB was a qualitative hypothesis: "networks learn low-order statistics first." The paper converts this into two precise, falsifiable Criteria (Section 2.2) that are independently testable and complementary. Criterion 1 (grafting) asks: if I surgically replace the means and covariances of a Pekinese dog with those of a goldfish, does the network call it a goldfish? Criterion 2 (deletion) asks: if I remove everything from a distribution except its means and covariances, does the network's accuracy survive?

The intellectual contribution here is not the specific optimal transport or maximum-entropy algorithms (those are in Section 3), but the diagnostic framing these criteria enable. Criterion 1 is an active intervention—it tests whether the network uses the grafted statistics for classification. Criterion 2 is a deletion test—it tests whether the network needs anything beyond the preserved statistics. Together they provide converging evidence: if grafting changes classification and deletion doesn't harm performance, the network is genuinely operating at that statistical order.

This is a significant methodological contribution because it gives the field a standardized way to probe any trained network's statistical sensitivity without modifying the training procedure. The framing is broadly applicable: future work can apply these same criteria (with the same optimal transport and maximum-entropy code, which the authors release) to new architectures, new datasets, and new modalities. The fact that both criteria yield consistent results across ConvNeXt V2, Swin Transformer V2, and RegNet-Y (Figure 3, Appendix C) is strong evidence that the DSB is not an artifact of a particular probing technique—it is robust to how you measure it.

The maximum-entropy principle in Criterion 2 is an especially elegant choice. By appealing to Jaynes (1957), the paper grounds the notion of "deleting higher-order information" in a principled information-theoretic framework rather than an ad-hoc truncation. The maximum-entropy distribution is the unique distribution that satisfies the moment constraints while introducing the minimum additional structure—it is the most conservative baseline for what the network could possibly be using. If a network's accuracy on maximum-entropy samples equals its accuracy on real data, there is literally no information in the real data beyond those moments that the network is exploiting.

Innovation 3: The N-Gram/Moment Equivalence as a Unifying Framework Across Modalities

The paper's extension of the DSB to discrete domains (Section 2.5, Theorems 2.1-2.2) is more than a technical convenience—it is a conceptual unification. Prior work on learning order in language models (Choshen et al., 2022) observed that neural LMs initially match n-gram model behavior, but this was an empirical observation without a theoretical link to the moment-based DSB framework used for images. The paper closes this gap by proving that token n-gram frequencies are exactly equivalent to the moments of one-hot encoded sequences, and that this equivalence is preserved under linear embedding maps (the matrix multiplication that converts one-hot vectors to continuous embedding vectors).

The significance of this equivalence is that it allows the same diagnostic framework (maximum-entropy sampling, testing at log-spaced training checkpoints, looking for non-monotonic U-shaped or double-descent curves) to be applied to language models without modification. A k-gram autoregressive model is precisely the maximum-entropy distribution matching all n-gram statistics up to order k. This means the language modeling experiments in Section 4 are not merely analogous to the vision experiments—they are the same experimental protocol applied to a different modality, made possible by the theoretical bridge.

Theorem 2.2 deepens this by showing that matching n-gram frequencies forces matching moments in the embedding space where the neural network actually operates. This is not obvious: two token sequences with the same bigram statistics could, in principle, map to different distributions in the continuous embedding space if the embedding matrix creates higher-order correlations that diverge. The theorem proves this doesn't happen—the linearity of the embedding map guarantees that equal n-gram statistics up to order k imply equal embedding-space moments up to order k. This is what makes the Taylor expansion motivation (Equation 2) applicable to language models: the network's loss depends on embedding-space moments, and those moments correspond to n-gram statistics.

The unification also provides a theoretical foundation for a phenomenon that Choshen et al. (2022) observed empirically but couldn't fully explain: why do neural LMs initially behave like k-gram models before diverging? The DSB provides the answer: because k-gram models capture the first k moments of the token distribution, and networks learn these moments in increasing order. The initial alignment is not a coincidence or a metaphor—it is a direct consequence of the moment structure of the data and the monotonic learning order.

Innovation 4: The Double Descent of Statistical Complexity in Language Models

The most surprising empirical finding in the paper is the "double descent" pattern observed in Pythia language models evaluated on low-order n-gram sequences (Figure 6). Unlike the vision experiments, where accuracy on low-order synthetic data follows a simple U-shape (rise then fall), the language model loss on unigram and bigram sequences first drops, then rises again (the first descent and ascent, matching the vision U-shape), but then drops a second time later in training. This second descent is substantial—the loss on unigram and bigram sequences reaches its global minimum late in training, not at the U-shape trough.

The paper's diagnosis of this phenomenon is that the second descent is caused by the emergence of in-context learning. Figure 7 shows that for early training steps (step 16, step 256), the loss on n-gram sequences is flat across token positions—the model treats each token independently, as expected if it's using only low-order statistics. But by step 8,000 and especially step 66,000 and step 143,000, the loss decreases as more tokens of the sequence are observed. This is the signature of in-context learning: the model is inferring the data-generating process (the n-gram model that produced the sequence) on the fly and using that inference to improve its predictions on later tokens.

This is a conceptually rich finding because it reveals an interplay between two different modes of statistical learning. The DSB describes the parametric learning—what the network's weights encode about the training distribution. But in-context learning is a non-parametric capability—the model can adapt to novel statistical structure within a single forward pass, without weight updates. The double descent occurs because these two modes have opposite temporal trajectories: parametric reliance on low-order statistics peaks early then fades (the U-shape), while in-context learning ability grows monotonically with training (the second descent). The net loss curve is the sum of these two effects.

The paper also notes a fascinating scale dependence: smaller Pythia models appear to resume the standard U-shaped pattern in the very late stages of training, while larger models sustain the in-context-learning-driven second descent. The authors speculate that larger models have greater representational capacity to retain the early-learned n-gram features while also developing in-context learning, whereas smaller models experience a form of "catastrophic forgetting" where the n-gram features are overwritten. If confirmed, this would add a new dimension to our understanding of how model scale affects the balance between parametric and in-context learning, and it aligns with broader observations that larger models are better at retaining diverse capabilities.

Innovation 5: Verifying That Low-Order Statistics Suffice for "Easy" Datasets

While not as flashy as the double-descent finding, the paper's comparative results across datasets with different complexity levels (Figure 2, Appendix C) make a conceptually important point: the DSB's U-shaped curve is most dramatic when the dataset cannot be solved using only low-order statistics. For MNIST and Fashion MNIST, where second-order maximum-entropy samples are still clearly recognizable as digits and clothing items (Figure 2), the accuracy curves show only a modest U-shape—networks achieve high accuracy on second-order samples and largely maintain it throughout training, because the low-order statistics genuinely contain enough information to achieve good performance. For CIFAR-10, where second-order samples look like colored noise (Figure 2), the U-shape is deep and pronounced—networks temporarily achieve better-than-random accuracy on low-order samples, then plummet as they incorporate the higher-order structure needed for genuine CIFAR-10 classification.

This is not just a robustness check—it's a validation of the DSB's explanatory scope. The DSB predicts that networks will learn whatever statistical order is sufficient to achieve low loss, and will stop "moving up" the order hierarchy once the data can be explained. For MNIST, the means and covariances are nearly sufficient, so the network never strongly relies on higher-order statistics and the U-shape is shallow. For CIFAR-10, the means and covariances capture color and blurry shapes but none of the fine-grained texture needed to distinguish a cat from a dog (compare the CIFAR-10 fake images in Figure 2—the airplane, bird, and deer classes produce almost indistinguishable noise patterns from second-order statistics alone), so the network must progress to higher orders, producing a deep U-shape. This dataset-dependence confirms that the DSB is not a rigid schedule but a sufficiency-driven progression: networks extract the simplest statistics that work, and move on only when those stop helping.

This finding also has practical implications for data augmentation and curriculum learning. If a dataset is "easy" in the sense that low-order statistics are highly discriminative (like MNIST), aggressive data augmentation that randomizes higher-order structure should not hurt—and indeed the paper's standard training uses RandAugment, random flips, and random crops, and the DSB pattern still holds. For "hard" datasets like CIFAR-10, the transition away from low-order statistics is where augmentation likely has its biggest impact, by making it harder for the network to overfit to spurious low-order correlations. The paper's framework provides a principled way to diagnose where a given dataset sits on this spectrum.

5. Experimental Analysis

Evaluation Methodology

• Datasets. The image classification experiments use CIFAR-10 (Krizhevsky et al., 2009), Fashion MNIST (Xiao et al., 2017), MNIST (LeCun et al., 1998), and SVHN (Netzer et al., 2011)—all at 32×32 resolution—plus a custom CIFARNet dataset consisting of 200K images at 64×64 resolution sampled from ImageNet-21K using 10 coarse-grained classes roughly matching CIFAR-10 (20K images per class; see Appendix B for synset details). For language modeling, the training corpus is the Pile (Gao et al., 2020) and n-gram frequencies are computed from Pythia's training data; evaluation sequences of length 2049 are sampled from unigram, bigram, trigram, and 4-gram autoregressive models fitted to these frequencies (3- and 4-gram models use a subset of the Pile for tractability).

• Base models. Computer vision experiments use ConvNeXt V2 (Woo et al., 2023), Swin Transformer V2 (Liu et al., 2022), and RegNet-Y (Radosavovic et al., 2020). ConvNeXt and Swin are swept across Atto, Femto, Pico, Nano, and Tiny sizes; Swin sizes below Tiny are constructed by copying the embedding dimension from the corresponding ConvNeXt size. RegNet-Y serves as an optimizer/schedule robustness check with SGD and momentum (no LR warmup). Language modeling experiments use the full Pythia suite (Biderman et al., 2023) from 14M to 12B parameters, for which log-spaced training checkpoints are publicly available, plus custom-trained Pythia 14M and 70M variants with extended learning rate warmup (14,300 steps) to isolate warmup effects.

• Metrics. For optimal transport experiments, the primary metric is classification accuracy computed with respect to the target class label—i.e., an image from source class A edited to match class B's statistics is evaluated against the ground-truth label for B, and accuracy above the 10% random baseline (for 10-class datasets) indicates the model is sensitive to the grafted statistics. For maximum-entropy sampling experiments in vision, both accuracy and cross-entropy loss are reported; in language modeling, the primary metric is cross-entropy loss in bits per byte (bpb) plus KL divergence between Pythia's token predictions and the n-gram model's predictions. The paper reports the Shannon entropy of the unigram distribution (2.89 bpb) and bigram distribution (2.04 bpb) as reference lower bounds for the loss curves in Figure 6.

• Baselines. The paper's experimental design does not use learned baselines in the conventional sense—there is no "prior method" against which the DSB is benchmarked. Instead, the baselines are the random classification baseline (10% for 10-class problems), which establishes the floor for optimal transport accuracy (values above this indicate the model is not merely guessing), and the validation set performance on real data, which establishes the ceiling that maximum-entropy samples could potentially match if the model relies only on the preserved statistics. For language modeling, the entropy of the n-gram models themselves serves as a natural baseline: Pythia's loss cannot go below this without in-context learning that infers the n-gram data-generating process within the sequence. The paper also implicitly uses the training step as a temporal baseline: the key comparisons are across training steps within the same model and probe type, not across different models or methods.

• Generation budget / compute accounting. There is no test-time compute budget parameter in this paper—the "budget" is the training step count, evaluated at log-spaced intervals. Vision models are trained for 2^16 steps with batch size 128 using AdamW (learning rate 10^−3, linear decay, 2000-step warmup; SVHN uses 10^−4 for convergence) and standard data augmentation (RandAugment, random horizontal flips, random crops). Language models use the publicly available Pythia checkpoints at their native training intervals. The synthetic data generation cost is accounted for separately in Appendix G: for CIFARNet (the most expensive case), second-order hypercube-constrained sampling costs ~65 seconds on an L40 GPU, while third-order matching requires ~36 hours on an A40 GPU due to the coskewness tensor optimization.

• Cross-validation / statistical protocol. For the Pythia language model experiments, the paper leverages multiple independent training runs to assess seed sensitivity: 10 seeds each for Pythia 14M, 70M, and 160M, and 5 seeds for Pythia 410M. Results across seeds are plotted as translucent lines in Figure 14 (Appendix C.6) with the claim that "random seed has very little effect on the learning curves." For vision experiments, individual model scales within an architecture family are plotted as translucent lines with the bold line showing the average across scales (Figure 3, Figure 5). There is no explicit cross-validation on held-out prompts or difficulty bins—the experimental design is fundamentally about within-training-trajectory comparisons rather than generalization to new data, and the synthetic probes themselves are constructed from the training set statistics.

Main Quantitative Results

The paper's results are organized by probing methodology (optimal transport for Criterion 1, maximum-entropy sampling for Criterion 2) across vision and language modalities, with difficulty-based analysis via dataset complexity rather than per-example difficulty estimation.

Optimal Transport (Criterion 1: Grafting) in Vision Models

Headline finding: All tested architectures classify OT-edited images according to their grafted statistics well above chance early in training, peak at intermediate training steps (roughly step 2^9 = 512), then decline, producing a characteristic U-shaped accuracy curve. The peak accuracy and the depth of the subsequent decline depend on the dataset and the order of statistics being grafted.

CIFAR-10 results (Figure 3, Appendix C.1):

• Second-order Gaussian OT achieves the highest peak accuracy among all interventions: ConvNeXt reaches over 40% accuracy with respect to the target class on CIFAR-10 at approximately step 2^10 (1024), compared to the 10% random baseline. This means that for over 40% of test images edited to match another class's mean and covariance, the model classifies them as the target class rather than the original source class.

• First-order CQN peaks earlier and lower: accuracy rises to roughly 20-25% around step 2^8–2^9, then declines. The fact that CQN-edited images (which preserve source-class correlations) are classified as the target class less often than Gaussian OT-edited images (which impose the target-class covariance structure) indicates that correlations between pixels contain class-relevant information that CQN leaves intact from the source class—the model can still "see through" the first-order edit to some degree because the spatial structure remains.

• First-order bounded shift (which changes only the per-image mean, matching the class-conditional mean approximately) shows similar or slightly lower accuracy than CQN, consistent with it being a weaker intervention on first-order statistics.

• Non-monotonicity is evident across all architectures: accuracy rises from near-random at initialization, peaks, then declines below the random baseline for some interventions (particularly first-order ones), indicating that late-training models actively reject the grafted-statistics classification and instead classify based on higher-order structure that contradicts the grafted low-order statistics.

• Validation set accuracy (the dark line in each plot) increases monotonically throughout training, serving as a control that the models are genuinely learning the real classification task. The contrast between the monotonically improving validation accuracy and the U-shaped synthetic-probe accuracy is the paper's central visual argument.

CIFARNet results (Figure 8, Appendix C.2): The 64×64 resolution dataset qualitatively mirrors CIFAR-10 but with somewhat lower peak accuracies on second-order samples. The paper hypothesizes that "networks more quickly learn to use higher-order statistics when the input has higher dimensionality," consistent with the DSB: larger input dimension means more opportunity for higher-order moments to carry discriminative information, reducing the window during which low-order statistics dominate.

Fashion MNIST and MNIST results (Figures 9-10, Appendices C.3-C.4): Both datasets exhibit "only a modest degree of non-monotonicity." The maximum-entropy synthetic images from second-order statistics for these datasets are still clearly recognizable (Figure 2), meaning low-order statistics genuinely suffice for good classification. Accuracy on synthetic probes rises and then plateaus or declines only slightly, because real validation performance is already largely achievable from first- and second-order information alone.

SVHN results (Figure 11, Appendix C.5): This dataset is an outlier: "none of the models ever exceed random baseline accuracy on 1st order synthetic images." The paper hypothesizes that the extreme diversity of colors, fonts, and background textures in SVHN makes first-order features non-discriminative—the per-pixel means across different digit classes are too similar to distinguish them, so even early-training models cannot classify based on first-order statistics alone.

Architecture and scale robustness: The U-shaped pattern is consistent across ConvNeXt V2, Swin Transformer V2, and RegNet-Y (optimized with SGD+momentum rather than AdamW), and across model scales from Atto to Tiny. The paper states: "model scale has a remarkably small effect on the learning curves," with the bold average line tightly tracking the individual translucent lines. This is evidence that the DSB is not architecture-specific or scale-dependent—it appears to be a general property of gradient-based training on these datasets.

Maximum-Entropy Sampling (Criterion 2: Deletion) in Vision Models

Headline finding: Accuracy on maximum-entropy synthetic data peaks earlier in training and at lower maximum values for lower-order matching (first-order before second-order), with the highest accuracy achieved on hypercube-constrained second-order samples. Across all datasets, some degree of U-shaped non-monotonicity is observed. Strikingly, for some datasets early in training, models achieve higher accuracy on independent pixel samples (ICS, first-order) than on the real validation set.

CIFAR-10 results (Figure 5, Appendix C.1):

• First-order ICS (matching per-pixel marginals but no correlations): accuracy peaks early at roughly 25-30% for ConvNeXt around step 2^7–2^8, then declines toward or below the random baseline. The paper notes the remarkable result that early-training accuracy on ICS samples can exceed real validation accuracy—the model is temporarily better at classifying these simplified, maximum-entropy images than the real complex images.

• First-order Conrad sampling (matching only per-pixel means, maximizing entropy subject to hypercube constraint): similar temporal pattern to ICS, with slightly lower peak accuracy, confirming that the per-pixel means (not the full marginal distribution) drive most of the first-order effect.

• Second-order Gaussian sampling (unconstrained support, matching means and covariances): accuracy peaks later (roughly step 2^9–2^10) and at higher values (roughly 40-50% depending on architecture) compared to first-order probes. The later peak and higher maximum confirm that second-order statistics become usable by the network after first-order statistics and provide more discriminative information.

• Second-order truncated normal (hypercube-constrained, matching means and covariances but with pixel values in [0,1]): achieves the highest peak accuracy among all probes—higher than unconstrained Gaussian sampling. This is important because it demonstrates that the Gaussian accuracy is not artificially depressed by out-of-range pixel values; constraining to [0,1] while matching the same moments improves performance, suggesting the network finds these more "image-like" samples easier to process.

• Third-order hypercube-constrained (matching means, covariances, and coskewness): accuracy peaks even later and higher than second-order, consistent with the monotonic-order prediction. However, the cost of generating these samples limits their use primarily to CIFAR-10 and CIFARNet.

• The ordering of peak accuracy across probe types follows the statistical order: first-order < second-order < third-order, and the timing of the peaks similarly follows this order (first-order peaks earliest, third-order peaks latest). This is the most direct empirical evidence for the DSB's core claim: networks learn to exploit statistics in increasing order.

Fashion MNIST and MNIST (Figures 9-10): The U-shape is shallow—accuracy on second-order samples rises and largely plateaus rather than declining substantially, because the second-order samples already look realistic (Figure 2). This validates the DSB's prediction that networks will not strongly move beyond low-order statistics when those statistics already capture the discriminative structure of the data.

CIFARNet (Figure 8): Maximum accuracies on second-order samples are somewhat lower than for CIFAR-10, consistent with the hypothesis that higher-dimensional inputs cause faster progression to higher-order statistics. The U-shape is still present but the peak is lower, suggesting the window of low-order reliance is narrower at higher resolution.

Loss curves: For maximum-entropy probes, the paper also reports cross-entropy loss (bottom panels in each figure). These show the inverse pattern to accuracy: loss on low-order probes reaches a minimum early in training, then rises later. The loss on the validation set monotonically decreases. The gap between the synthetic-probe loss minimum and the validation loss at that same step quantifies how much "extra" information the real data contains beyond the preserved moments—a gap that widens as training proceeds and the network becomes sensitive to higher-order statistics.

Language Modeling: N-Gram Sequence Loss and Double Descent

Headline finding: Pythia language models exhibit a "double descent" pattern on unigram and bigram sequence loss: loss drops to a trough between step 2^6 (64) and 2^8 (256), rises until approximately step 2^10 (1024), then drops again later in training. This contrasts with the vision experiments' simple U-shape. The paper attributes the second descent to the emergence of in-context learning.

Aggregate loss curves (Figure 6, Appendix C.6):

• Unigram sequence loss: The unigram entropy is 2.89 bpb. Pythia 12B's unigram sequence loss drops from ~5 bpb at initialization to a minimum of roughly 2.5-2.6 bpb around step 2^7–2^8, rises to roughly 3.0-3.2 bpb around step 2^10–2^12, then declines again to ~2.0 bpb by the end of training. The final loss going below the unigram entropy (2.89 bpb) is only possible through in-context learning—the model cannot achieve <2.89 bpb by parametric knowledge alone, because the unigram model's predictions are maximum-entropy given unigram frequencies; beating this lower bound requires inferring from context that the sequence is unigram-generated and adjusting predictions accordingly.

• Bigram sequence loss: The bigram entropy is 2.04 bpb. The double-descent pattern is similar: loss drops from ~5 bpb to a minimum of ~2.0-2.2 bpb around step 2^7, rises to ~2.5-2.8 bpb around step 2^10–2^12, then declines to ~1.5-1.7 bpb by the end of training for larger models. Again, going below 2.04 bpb requires in-context learning.

• Trigram and 4-gram sequence loss: These do not exhibit the double-descent non-monotonicity—loss monotonically decreases throughout training after an initial drop. However, the paper notes that the 3-gram loss "plateaus at a point consistent with the non-monotonicity observed in 1- and 2-grams, suggesting that the effect could be present to a lesser extent."

• Model scale dependence: Smaller models show more pronounced U-shaped re-ascent in the late-training regime for unigram and bigram sequences, while larger models sustain the second descent. The paper speculates: "we speculate that this behavior may arise from a form of 'catastrophic forgetting', in which all models initially learn low-order n-gram statistics, which are eventually eclipsed by more sophisticated features. Larger models have greater representational capacity, and so are better able to retain these early n-gram features." This interpretation is supported by the observation that Pythia 14M's unigram loss rises substantially in late training (from ~2.8 bpb back to ~3.5+ bpb), while Pythia 12B's unigram loss stays near its global minimum.

• Seed robustness (Figure 14): Across 9 seeds for Pythia 14M, 70M, and 160M, and 4 seeds for Pythia 410M, the n-gram loss curves are "very consistent," with translucent individual-seed lines tightly clustered around the mean. This indicates the double-descent pattern is not a fluke of a particular random initialization.

In-context learning evidence (Figure 7): The paper follows Kaplan et al. (2020) in operationalizing in-context learning as decreasing loss at increasing token indices within a sequence. Analysis of Pythia 12B at six training checkpoints reveals:

• At step 16 and step 256: loss on unigram and bigram sequences is approximately uniform across token positions, indicating no in-context learning. The model treats each token independently based on its parametric knowledge.

• At step 1,000: some in-context learning of unigram sequences begins to appear—loss decreases slightly across token positions.

• At step 8,000: clear in-context learning of bigram sequences emerges, with loss decreasing substantially over successive tokens. This step is approximately where the aggregate bigram loss curve begins its second descent.

• At step 66,000 and step 143,000: strong in-context learning is present for both unigram and bigram sequences, with loss dropping steeply across early token positions before leveling off.

The temporal alignment between the emergence of in-context learning (Figure 7) and the second descent in aggregate loss (Figure 6) provides correlational evidence for the paper's hypothesis that in-context learning drives the double descent. The paper characterizes this as "correlational evidence" rather than causal proof, acknowledging the limitation.

KL divergence analysis (Figure 12, bottom panels): The KL divergence between the n-gram model's token predictions and Pythia's predictions mirrors the loss curves. For unigrams, KL divergence drops sharply early, rises, then drops again—the second drop indicating that Pythia's predictions are converging toward the unigram model's predictions (but going beyond them, as evidenced by lower loss). For bigrams, the pattern is qualitatively similar.

Learning rate warmup effect (Figure 13): Custom Pythia 14M and 70M models with extended warmup (14,300 steps vs. the default fast warmup) show similar double-descent patterns, with the paper noting that "lengthening the LR warmup period did not consistently affect their overall shape." This rules out the concern that the double descent is an artifact of the specific learning rate schedule.

Ablation Studies and Robustness Checks

Architecture family: The U-shaped accuracy curves for optimal transport and maximum-entropy sampling are replicated across ConvNeXt V2 (CNN-based), Swin Transformer V2 (transformer-based), and RegNet-Y (CNN, SGD optimizer), shown in Figure 3, Figure 5, and all panels of Appendix C. The consistency across these architecturally diverse models is the paper's strongest evidence that the DSB is not architecture-specific.

Model scale within architecture: Within each architecture, model sizes from Atto (~3.7M parameters) to Tiny (~49M parameters) produce nearly overlapping translucent curves in Figures 3 and 5, with the paper stating "model scale has a remarkably small effect on the learning curves." This suggests the DSB is a property of the learning dynamics rather than a capacity constraint—even large models that could hypothetically learn higher-order statistics from the start still follow the monotonic order.

Optimizer and learning rate schedule: RegNet-Y trained with SGD+momentum (no LR warmup, unlike AdamW with warmup) reproduces the U-shaped curves, confirming the DSB is not an artifact of adaptive optimization or warmup schedules (Appendix C.1-C.5).

Dataset difficulty / complexity: The degree of non-monotonicity varies systematically with dataset complexity: MNIST and Fashion MNIST (where second-order samples are realistic, Figure 2) show shallow U-shapes; CIFAR-10 and CIFARNet show deep, pronounced U-shapes; SVHN shows no above-chance accuracy on first-order probes at any point in training. This is not presented as a formal ablation but serves as a natural robustness check: the DSB predicts and the experiments confirm that the progression through statistical orders is conditioned on the informativeness of lower-order statistics.

First-order method variants (CQN vs. bounded shift vs. Conrad vs. ICS): The paper includes multiple first-order methods that differ in what exactly is matched (full marginals via CQN, means via bounded shift, means via Conrad maximum-entropy, full marginals via ICS independent sampling). All produce qualitatively similar temporal patterns, with the ordering of peak accuracy (ICS ~ CQN > Conrad ~ bounded shift) reflecting the amount of first-order information preserved. This suggests the DSB's first-order effect is robust to how exactly "first-order statistics" are operationalized.

Second-order method variants (Gaussian vs. truncated normal): For CIFAR-10 and CIFARNet, the hypercube-constrained second-order samples produce higher peak accuracy than unconstrained Gaussian samples, confirming that the U-shape on Gaussians is not driven by out-of-distribution pixel values. The hypercube constraint brings samples closer to the real data manifold, making them easier for the network to process—but the temporal pattern (peak timing, subsequent decline) is preserved.

Third-order hypercube-constrained sampling: The availability of third-order probes for CIFAR-10 and CIFARNet (matching means, covariances, and coskewness, with hypercube constraint) provides a direct test of the monotonic-order prediction: third-order probes should peak later and at higher accuracy than second-order probes. The results in Figure 5 and Appendix C confirm this, though the computational cost limits the scope.

Language model seed robustness: Multiple random seeds for Pythia 14M, 70M, 160M (9 seeds each) and 410M (5 seeds) produce highly consistent n-gram loss curves (Figure 14), ruling out seed-specific artifacts in the double-descent pattern.

Language model learning rate warmup: Extended warmup Pythia variants (14,300 steps vs. default) show the same double-descent pattern (Figure 13), ruling out warmup-duration artifacts.

N-gram order comparison in language models: The paper evaluates 1-, 2-, 3-, and 4-gram sequences, finding that the double-descent is most pronounced for 1- and 2-grams, present but subtle for 3-grams, and absent for 4-grams. The ordering of loss minima (1-gram reaches minimum first, 4-gram last) is consistent with the DSB, though the interpretation is complicated by the in-context learning effect.

Negative result: in-context learning as a confound for higher-order n-gram probes. The paper does not find clear double-descent for 3- and 4-gram sequences. This could be because in-context learning of higher-order n-gram structure is harder and emerges later (or not at all), or because the parametric learning of these higher-order statistics dominates throughout training. The paper does not resolve this ambiguity, noting only that "the 3-gram loss plateaus at a point consistent with the non-monotonicity observed in 1- and 2-grams."

Critical Assessment

The paper makes one central empirical claim: neural networks learn statistics of increasing complexity (low-order moments first, higher-order moments later), and this progression is visible in both vision and language models through a characteristic U-shaped or double-descent accuracy/loss curve on synthetic statistical probes. The evidence supporting this claim is extensive but comes with specific boundary conditions and methodological caveats that the paper is largely transparent about.

What the experiments demonstrate vs. what they do not:

The experiments convincingly demonstrate that at specific intermediate training checkpoints, networks achieve higher accuracy (or lower loss) on synthetic data matching only low-order statistics than at earlier or later checkpoints. This is the U-shape that the paper's framework explicitly predicts. The fact that this pattern appears across five vision datasets, three architectures, multiple optimizer configurations, and model scales from Atto to Tiny (with "remarkably small" scale effects) makes a strong case that this is a general phenomenon, not a peculiarity of a particular setup.

However, what the experiments demonstrate is correlational evidence for the DSB's temporal ordering claim, not mechanistic evidence. The paper shows that the network's behavior is consistent with learning first-order before second-order before third-order statistics, but it does not show why this ordering occurs. The Taylor expansion motivation (Equation 2) is an argument for plausibility, not a formal proof that gradient descent on neural network loss functions necessarily produces monotonic moment learning. The paper does not analyze the actual derivatives during training, does not measure the correlation between derivatives and moments at different training steps, and does not rule out alternative explanations for why the U-shape occurs (e.g., it could be related to the loss landscape geometry rather than moment ordering per se).

Specific strengths of the experimental design:

1. The inversion of Refinetti et al.'s paradigm (training on real data, probing with synthetic data) is genuinely stronger evidence than training on synthetic approximations. It shows that the DSB operates even when the network has access to all statistical orders from the start.

2. The two-criterion design (grafting AND deletion) provides converging evidence. That both criteria produce consistent U-shaped patterns across all datasets substantially reduces the risk that the results are an artifact of a particular probing technique.

3. The dataset-dependence of the effect (strong U-shape on CIFAR-10, weak on MNIST, absent on SVHN first-order) is not just consistent with the DSB—it is a necessary pattern if the DSB is correct. If low-order statistics suffice (MNIST), the U-shape should be shallow. If low-order statistics are non-discriminative (SVHN first-order), no peak should occur. The experiments confirm both predictions.

4. The language modeling extension through Theorems 2.1-2.2 is a non-trivial theoretical bridge that enables cross-modal comparison. The double-descent finding in Pythia, with the in-context learning diagnosis via positional loss analysis (Figure 7), adds a dimension to the DSB that the vision experiments alone would not reveal.

Specific weaknesses and missing experiments:

1. The causal link to the Taylor expansion is untested. The paper motivates the DSB through Equation 2 but never measures any of the quantities in that equation during actual training—no derivative norms, no moment-contribution decomposition, no tracking of which Taylor terms dominate the loss at different training stages. The Taylor expansion serves as a conceptual metaphor but is not empirically validated as the mechanism behind the U-shapes.

2. Single language model family (Pythia). All language modeling results are on Pythia models trained on the Pile. There is no replication on other model families (e.g., GPT-2, LLaMA, OPT) or other corpora, limiting the claim that the double descent generalizes across language models. The theoretical bridge (Theorems 2.1-2.2) is model-agnostic, so replication on other architectures would substantially strengthen the paper.

3. Skip-grams are omitted from language model experiments. The paper's formal definition of n-gram statistics (Definition F.1) includes skip-grams, but the experiments use only contiguous n-grams "for tractability reasons." This means the synthetic language data does not actually match all moments up to order k—it matches only a subset (contiguous n-grams at absolute positions). Whether including skip-grams would change the U-shape or double-descent pattern is unknown.

4. No causal intervention on in-context learning. The paper identifies in-context learning as the driver of the second descent based on correlational timing evidence (comparing Figures 6 and 7). There is no experiment that directly tests this claim—for instance, by evaluating models where in-context learning is ablated (e.g., through attention pattern manipulation) or by training models known to lack in-context learning and checking whether the second descent disappears. The in-context learning interpretation is plausible but unproven.

5. Limited third-order probing in vision. Third-order hypercube-constrained samples are computationally expensive and were generated only for CIFAR-10 and CIFARNet. Expanding to MNIST and Fashion MNIST would test the DSB's prediction that third-order probes should add little beyond second-order for datasets where second-order statistics already suffice—a stronger test than the current indirect evidence.

6. No probing of fourth-order or higher statistics. The paper's evidence for "statistics of increasing complexity" stops at order 3 in vision and order 4 (contiguous n-grams) in language. Whether the monotonic progression continues to higher orders, and at what point it saturates, is untested.

7. The computational cost of difficulty estimation is absent (because there IS no difficulty estimation). This is a strength in one sense—the paper avoids the expensive oracle-difficulty estimation problem that the reference paper struggled with. But it also means that the paper provides no per-example analysis: are there some CIFAR-10 images where the U-shape is much more pronounced than others? Is there heterogeneity within a class? The aggregate curves in Figure 3 average over all test images; per-image difficulty analysis could reveal whether the DSB operates uniformly or only on a subset of examples.

8. No test of whether the DSB can be deliberately counteracted. A natural follow-up experiment would be to train with a loss that explicitly penalizes reliance on low-order moments (e.g., by adversarially perturbing the first two moments during training) and check whether the U-shape disappears. This would test whether the DSB is a soft bias that can be overridden or a hard constraint on learning dynamics.

9. No direct measurement of moment usage in the trained network. The criteria test sensitivity to interventions on moments, but they don't directly measure whether the network's internal representations encode those moments. A complementary analysis using probing classifiers or representational similarity metrics could strengthen the claim that the network is genuinely "using" the moments rather than merely being sensitive to them.

Conditional nature of the claims:

The DSB holds most cleanly when the following conditions are met: (a) the dataset requires higher-order statistics for good performance (CIFAR-10, not MNIST), (b) the low-order statistics are at least somewhat discriminative (CIFAR-10, not SVHN first-order), and (c) the model is evaluated during training, not just at convergence (the entire U-shape would be invisible if evaluating only the final checkpoint). For language models, the additional condition is that in-context learning must be separately accounted for—the raw loss curve conflates parametric DSB effects with in-context learning effects, and the double descent is only interpretable when these are distinguished (as in Figure 7).

The paper's most significant claim—that the DSB is a general property of neural network training—is supported with evidence that is broad (many datasets, architectures, scales) but not deep (no mechanistic verification, limited higher-order probing, single language model family). The consistency of results across the tested conditions is impressive, but the claim would be considerably strengthened by the missing experiments noted above. The framework and criteria the paper introduces are arguably its most lasting contribution, providing a standardized methodology that future work can apply to new models, datasets, and modalities to further test the generality and limits of the DSB.

6. Limitations and Trade-offs

The Taylor Expansion Motivation Is Not Empirically Validated as a Causal Mechanism

The assumption or constraint. The paper's entire theoretical framing rests on the Taylor expansion of the expected loss (Equation 2), which expresses $\mathbb{E}[L(x)]$ as an infinite sum over central moments of the data distribution weighted by normalized derivatives of the loss at the mean. The paper argues that the factorial decay of $1/\alpha!$ coefficients at initialization causes low-order terms to dominate, and that gradient descent should reduce these larger terms first, producing monotonic learning of moments. The paper explicitly acknowledges this is a motivation, not a proof:

"This connection provides some motivation for the DSB. Specifically, if during training, a network's loss is well approximated by the first k terms of its Taylor expansion, then the model should only be sensitive to statistics up to order k, and we argue that earlier terms of the expansion will generally become relevant before later terms." (Section 1)

"Another argument for monotonicity is that earlier terms account for factorially more of the loss at initialization, and are thus plausibly higher-priority targets for gradient descent, until the optimizer is no longer able to easily reduce the loss further by better matching the associated statistics and moves on to higher-order terms." (Section 2.1, footnote 6)

The paper also acknowledges that ReLU networks are not analytic, so the Taylor expansion, strictly speaking, does not apply:

"Famously, the ReLU activation function is not analytic, but it is possible to construct arbitrarily close approximations to ReLU that are analytic" (Section 2, footnote 2)

The consequence. The paper never measures any of the quantities in Equation 2 during actual training. It does not track the magnitude of derivatives $(\partial^\alpha L)(\mu)$ at different training steps, does not decompose the loss into contributions from different moment orders, and does not verify that the factorial coefficient argument holds for real networks trained with SGD/AdamW on real data. As a result, the U-shaped accuracy curves that constitute the paper's primary evidence for the DSB are correlational—they are consistent with the DSB hypothesis but do not rule out alternative explanations that have nothing to do with moment ordering. For example, the U-shape could arise from the geometry of the loss landscape (early-training networks occupy a flat, low-accuracy basin where they are easily fooled by simple perturbations; as training sharpens the loss landscape, the network becomes more sensitive to subtle features, producing the observed decline in probe accuracy). Without direct measurement of the Taylor expansion terms, the paper cannot distinguish between the DSB as a mechanism (gradient descent prioritizes low-order moments) and the DSB as a description of one observable consequence that could have other causes.

What evidence exists in the paper. There is none. The Taylor expansion appears in Section 2.1 and is then never revisited in any empirical analysis. No figure, table, or experiment measures derivative magnitudes, moment contributions, or Taylor term correlations. The evidence for the DSB is entirely behavioral: networks behave as if they learn moments in order, but the paper does not demonstrate that they do learn moments in order in the mechanistic sense that the Taylor expansion would imply.

Mitigation status. The paper does not attempt to close this gap. The Taylor expansion is presented as a framing device, and the authors treat the behavioral evidence (the U-shaped curves) as sufficient to establish the DSB's validity, without requiring mechanistic verification. Future work that directly measures the evolution of loss derivatives during training—for instance, by computing the correlation between $(\partial^\alpha L)(\mu)$ and $\mathbb{E}[(x-\mu)^\alpha]$ at different training steps for different orders $|\alpha|$—would substantially strengthen the DSB's theoretical basis.

---

Difficulty Estimation Cost Is Absent (and Perhaps Unnecessary), but Per-Example Variation Is Unexplored

The assumption or constraint. The paper evaluates all synthetic probes in aggregate, reporting average accuracy or loss across all test images or all sampled sequences. There is no per-example analysis of difficulty: no measurement of which CIFAR-10 images produce the strongest U-shaped response, no analysis of whether the DSB operates uniformly across all images or is driven by a subset, no attempt to characterize why some images might be more susceptible to first-order or second-order grafting than others. This is a deliberate design choice—the paper focuses on aggregate dynamics—but it leaves open a significant question about the heterogeneity of the DSB effect.

Unlike the reference paper (which explicitly studied difficulty-conditioned compute allocation and whose main contribution was a difficulty estimator), this paper has no difficulty estimation framework and does not frame its results in difficulty-dependent terms. The closest analog is the dataset-level comparison between MNIST/Fashion MNIST (where the U-shape is shallow because low-order statistics suffice) and CIFAR-10/CIFARNet (where the U-shape is deep because higher-order statistics are needed). But this is cross-dataset, not within-dataset.

The consequence. Without per-example analysis, we cannot determine whether the DSB's U-shaped aggregate curve reflects a uniform tendency (every image goes through the same statistical progression, just at different rates) or a compositional effect (some images are highly susceptible to low-order probes while others are never fooled). If the U-shape is driven by a small subset of "highly gatherable" images, the DSB may be a much weaker and more contingent phenomenon than the aggregate curves suggest. Conversely, if the U-shape is uniform across images, the DSB is a robust, universal learning dynamic. The paper provides no evidence either way.

This also has practical consequences for anyone trying to apply the DSB diagnostically. If the synthetic probes only work well on certain types of images (e.g., images with uniform backgrounds where mean grafting produces a strong signal, vs. cluttered images where it is imperceptible), then the probe methodology is not a general-purpose tool for assessing a model's statistical reliance—its effectiveness depends on which images you happen to test. A practitioner would want to know whether the probes are reliable across their specific data distribution.

What evidence exists in the paper. Essentially none. All figures (3, 5, 6, 8-11) report aggregate accuracy or loss. The error bars (translucent lines for individual model scales in Figures 3 and 5; translucent lines for individual seeds in Figure 14) capture variation across models and random seeds, not across examples within a dataset. Figure 2 provides anecdotal qualitative evidence that maximum-entropy samples vary dramatically in recognizability across datasets, but no per-image analysis is presented for synthetic probes within a single dataset. The paper never computes a per-image difficulty score (e.g., the accuracy of CQN-edited images binned by some image property) or performs the kind of difficulty-quintile breakdown that is common in related work.

Mitigation status. There is no attempt to address this. The paper's experimental framework is built to answer the question "does the DSB exist on average?" rather than "when and for which examples does the DSB operate?" This is a legitimate scoping decision, but it means the results are coarser than they could be. Future work that bins test images by properties relevant to each statistical order (e.g., images with high vs. low spatial autocorrelation, images with uniform vs. textured backgrounds, images where the class-conditional mean is distinctive vs. ambiguous) would reveal whether the DSB is a uniform or heterogeneous phenomenon. The language modeling experiments face the same issue—there is no analysis of whether the double descent is uniform across all n-gram sequences or is driven by sequences with particular statistical properties.

---

Single Language Model Family and Training Corpus

The assumption or constraint. All language modeling experiments use the Pythia model suite (Biderman et al., 2023) trained on the Pile (Gao et al., 2020). The paper does not evaluate any other language model architecture (e.g., LLaMA, GPT-2, OPT) or any other training corpus. While the paper validates the DSB across three vision architectures (ConvNeXt V2, Swin Transformer V2, RegNet-Y) and five vision datasets, the language modeling evidence comes from a single model family trained on a single dataset. The authors do not explicitly frame this as a limitation—they present Pythia as the natural choice because of its publicly available log-spaced checkpoints—but the asymmetry with the vision experiments is notable.

The consequence. The central contribution of Section 2.5 (Theorems 2.1 and 2.2) is a theoretical claim about any neural language model that uses embeddings: the equivalence between n-gram frequencies and embedding-space moments holds for any model, and therefore the DSB should be detectable in any language model through the maximum-entropy n-gram probing methodology. The theorems are model-agnostic, so they predict the double-descent pattern should appear in all autoregressive language models trained with gradient descent. But the paper tests this prediction on exactly one model family. We cannot distinguish between the following possibilities:

• The double descent is a universal property of language model training (as the theory predicts), and would replicate on LLaMA, GPT-2, Falcon, etc.
• The double descent is specific to Pythia's architecture, training data, or training recipe (e.g., the specific learning rate schedule, batch size, or data ordering used by Biderman et al., 2023).
• The double descent generalizes to other models but only when evaluated on n-gram distributions fitted to their training data, not the Pile.

The Pile itself has specific properties—it is a curated mixture of diverse text sources (Wikipedia, books, code, web text) with particular n-gram distributions. A corpus with different statistical structure (e.g., a purely code-based corpus, or a single-domain corpus with very different n-gram diversity) might produce a different temporal pattern in the n-gram loss curves, even under the same model architecture.

What evidence exists in the paper. The language modeling experiments are extensive within the Pythia family: all model sizes from 14M to 12B, multiple random seeds (9 each for 14M/70M/160M, 5 for 410M), custom warmup variants (Figure 13), and positional loss analysis for Pythia 12B (Figure 7). The paper also includes KL divergence measurements (Figure 12, bottom panels) and n-gram orders from 1 through 4. Within this scope, the double descent is robustly replicated. But there is zero evidence about whether it appears in non-Pythia models.

Mitigation status. No attempt. The paper does not discuss this as a limitation or call for replication on other model families. Given that the vision experiments carefully validate across architectures, optimizers, and datasets, the language experiments are notably narrower. The paper's public release of code and the straightforward nature of the probing methodology (computing n-gram frequencies and evaluating existing checkpoints) makes replication on other models relatively easy for future work, but the paper itself provides no such replication.

---

The In-Context Learning Interpretation of the Double Descent Is Correlational, Not Causal

The assumption or constraint. The paper attributes the second descent in unigram and bigram sequence loss to the emergence of in-context learning. The evidence for this is the temporal alignment between two observations: (1) aggregate unigram/bigram loss reaches its trough, rises, then drops again (Figure 6), and (2) the positional loss profile changes from flat (no in-context learning) to downward-sloping (in-context learning present) at a training step that approximately coincides with the second descent (Figure 7). The paper characterizes this cautiously:

"We hypothesize that this is caused by the development of in-context learning sufficient to improve n-gram sequence predictions. We observe correlational evidence in the n-gram sequence loss over increasing token indices and training steps in Pythia 12B (Fig. 7), where in-context learning seems to emerge in the same training step where the non-monotonic regime ends." (Section 4.1)

The paper does not perform any intervention that would establish causality: no ablation of in-context learning, no comparison to models incapable of in-context learning, and no demonstration that the second descent disappears when in-context learning is prevented.

The consequence. The double descent is the paper's most surprising and potentially significant finding—it suggests an interplay between parametric simplicity bias and non-parametric in-context learning that has not been documented before. If the causal link is real, it means that the DSB's U-shaped pattern in language models is "interrupted" by in-context learning, and that the two phenomena have opposite effects on n-gram sequence loss (the DSB pushes loss up as the model moves beyond low-order statistics; in-context learning pushes loss down as the model infers the n-gram data-generating process on the fly). This would be a conceptually rich finding.

However, there are plausible alternative explanations for the second descent that do not involve in-context learning:

• The later training stages might involve increased representation capacity for low-order statistics—as the model grows its overall representational power, it may re-learn to exploit low-order statistics in a way that is compatible with its more sophisticated higher-order features, producing a second drop in n-gram loss without any in-context learning. This would be a parametric, not non-parametric, phenomenon.
• The loss curve on n-gram sequences might simply be noisy or sensitive to the specific checkpoint schedule, and the apparent second descent could be an artifact of how the Pythia checkpoints are spaced. The paper's seed analysis (Figure 14) partially addresses this by showing consistency across runs, but if there is a systematic bias in the checkpoint schedule or the evaluation protocol, it would affect all seeds similarly.
• The model might be overfitting to the n-gram probes during evaluation, even though the probes are not part of training—the sequences are 4096 tokens long, and later checkpoints might have better internal representations for processing long sequences regardless of content.

Without causal evidence, the in-context learning interpretation remains an intriguing hypothesis supported by temporal correlation, but not proven.

What evidence exists in the paper. Two pieces: Figure 6 (aggregate loss curves showing double descent) and Figure 7 (positional loss analysis showing in-context learning). The paper also notes that the Pythia 12B loss on unigram sequences drops below the unigram entropy (2.89 bpb) in late training, which is only possible through in-context learning—parametric knowledge alone cannot beat the entropy of the data-generating process. This is stronger evidence for the presence of in-context learning, but still does not demonstrate that in-context learning causes the second descent rather than simply coinciding with it.

Mitigation status. The paper does not attempt to establish causality. No experiment intervenes on in-context learning (e.g., by truncating sequence length to prevent the model from accumulating sufficient context, by evaluating on shuffled n-gram sequences that break the Markov structure, or by using models trained with attention modifications that disable in-context learning). The paper acknowledges the correlational nature of the evidence but does not flag the absence of causal verification as a limitation to be addressed in future work. This is a gap—the double descent is sufficiently novel that confirming its mechanism would be high-priority for building on this paper's results.

---

Skip-Grams Are Included in the Theory but Excluded from Experiments

The assumption or constraint. The paper's formal definition of n-gram statistics (Definition F.1, Appendix F) explicitly includes skip-grams—an n-gram where tokens appear at specified absolute positions, which may be non-contiguous. For example, "the _ dog" (where _ matches any token) is a 2-gram with a gap. This definition is required for Theorem 2.1 to hold in its full generality: moments of the one-hot encoding correspond to conjunctions of (position, token) pairs, which can involve arbitrary subsets of positions, not just contiguous blocks.

However, the language modeling experiments use only contiguous n-gram models: a unigram model samples each token independently, a bigram model conditions each token on the immediately preceding token, a trigram model on the two preceding tokens, etc. The authors acknowledge this explicitly:

"While we include skip-grams (e.g. the _ dog) in our formal definition of n-gram frequency (Def. F.1), we do not include them in these experiments for tractability reasons: they would greatly increase the memory and storage requirements of maximum entropy sampling. We hope to explore the effect of skip-gram statistics in future work." (Section 4)

The consequence. The maximum-entropy n-gram language models used in the experiments do not match all moments up to order $k$ of the token distribution—they match only a subset of the order-$k$ moments, specifically those corresponding to contiguous token sequences at adjacent positions. Moments involving non-adjacent positions (e.g., the co-occurrence of a specific token at position 5 with another specific token at position 20) are not preserved by a bigram model, even though they are second-order moments of the one-hot encoding.

This means the language modeling experiments do not cleanly test the DSB's prediction that networks learn all order-$k$ statistics before order-$(k+1)$ statistics. Instead, they test a weaker claim: that networks first learn to model contiguous n-gram structure (unigram before bigram before trigram) before moving to higher-order contiguous structure. The theoretical bridge (Theorems 2.1-2.2) is built for the full moment definition including skip-grams, but the experiments test a restricted subset. The gap between the theory and the experiments is larger than the paper suggests.

Furthermore, the omission of skip-grams introduces an asymmetry with the vision experiments. In the vision case, "second-order" samples match the full covariance matrix, including all cross-pixel correlations (not just correlations between spatially adjacent pixels). The Gaussian optimal transport map $A$ in Equation 3 involves the full $d \times d$ covariance matrix, which includes long-range correlations between distant pixels. The language modeling probes, by using only contiguous bigrams, do not capture long-range token co-occurrence patterns that the network might be using. If Pythia models learn to exploit long-range token correlations early in training (which is plausible—certain function words like "the" and "of" have long-range co-occurrence patterns), the contiguous-bigram probes would fail to detect this, and the double-descent pattern might look different if full second-order moments were matched.

What evidence exists in the paper. None. There is no experiment comparing contiguous-only vs. skip-gram-inclusive n-gram probing, no analysis of how much additional information skip-grams would capture about the token distribution, and no measurement of whether Pythia models are sensitive to long-range n-gram statistics during early training. The KL divergence measurements (Figure 12) compare Pythia to the contiguous n-gram models only.

Mitigation status. Flagged as future work, not addressed. The tractability issue is genuine—a maximum-entropy model matching all n-gram statistics up to order $k$ with skip-grams would require conditioning on arbitrary subsets of previous tokens, which is combinatorially explosive and cannot be represented as a simple $k$-gram transition table. But the paper does not discuss whether a partial middle ground exists (e.g., including a subset of important skip-grams, or using a factorized approximation) that would bring the experiments closer to the theory without full combinatorial cost.

---

Higher-Order Probing Beyond Third Order Is Feasibility-Constrained, Leaving the DSB's Monotonicity Claim Incompletely Tested

The assumption or constraint. The paper's central empirical claim is that networks learn statistics in monotonically increasing order: first-order, then second-order, then third-order, and by implication, fourth-order, fifth-order, and beyond. However, the experiments only probe up to order 3 in vision (coskewness, via hypercube-constrained gradient optimization) and order 4 in language (contiguous 4-grams). Higher-order probes are computationally infeasible because the size of the moment tensors grows as $O(d^{\text{order}})$, where $d$ is the input dimension. The paper is transparent about this:

"Additionally, the maximum entropy third order hypercube-constrained sampling is much more expensive than the second order methods, since the size of the statistic tensor grows as O(d^order). This means the coskewness tensor for CIFARNet images has dimensions 12288 × 12288 × 12288. This would require nearly eight terabytes to store in full precision, which exceeds the memory capacity of our computing hardware by a significant degree." (Appendix G)

Even the third-order probes required substantial engineering (slice-based optimization to avoid materializing the full coskewness tensor). Fourth-order probes (cokurtosis) would require a tensor of size $d^4$, which is entirely out of reach—for CIFARNet with $d = 12288$, this would be ~$2.3 \times 10^{16}$ entries.

The consequence. The paper can only directly verify the monotonic learning order for the transition from order 1 to order 2, and from order 2 to order 3. Whether the progression continues to orders 4, 5, and beyond is an extrapolation, not an empirical finding. Several patterns could break the monotonicity at higher orders:

• Saturation: the network might stop progressing through statistical orders once it achieves sufficient accuracy on the training data, even if higher-order statistics exist and could be exploited. CIFAR-10 test accuracy for ConvNeXt V2 Nano is ~90%+ by late training (visible in the validation curves in Figure 3). If the network can achieve 95% training accuracy using statistics up to order 3, it might never learn to use order-4 statistics at all—the DSB would saturate rather than continuing indefinitely.

• Non-monotonic skipping: the network might learn order-3 statistics, then skip order 4 and learn order 5, because order-4 statistics of natural images might be less informative than order-5 statistics (e.g., for certain texture types). The Taylor expansion argument assumes a smooth, factorial decay in initial contribution, but if some higher-order moments are intrinsically more informative than intermediate-order ones, the learned order might not be strictly monotonic.

• Capacity bottlenecks: smaller models might simply lack the capacity to represent higher-order statistical dependencies, meaning the DSB progression would stop at a model-size-dependent ceiling rather than continuing through all orders. The paper's finding that model scale has "remarkably small effect on the learning curves" (Section 3.3) applies to the tested range (Atto to Tiny, roughly 4M to 49M parameters for ConvNeXt) and the tested orders (1-3). It does not guarantee that a 1M-parameter model would ever learn third-order statistics, or that a 1B-parameter model would learn tenth-order statistics.

What evidence exists in the paper. The ordering from first to second to third is consistent across CIFAR-10 and CIFARNet for the maximum-entropy probes (Figure 5, Appendix C.1-C.2): first-order ICS peaks earliest and lowest, second-order Gaussian peaks later and higher, third-order hypercube-constrained peaks latest and highest. This supports the monotonicity claim for the first three orders. But the paper provides no evidence for order 4 or beyond, and the computational discussion in Appendix G makes clear that direct verification at order 4 is infeasible with current methods.

Mitigation status. The paper acknowledges the computational infeasibility of higher-order probes but does not treat the missing orders as a limitation of the DSB claim. The authors frame the third-order results as confirmation of the monotonic progression and implicitly extrapolate to higher orders. The paper does not suggest an alternative methodology for probing higher-order statistics that would bypass the $O(d^{\text{order}})$ scaling (e.g., random projections of moment tensors, or probing specific subsets of higher-order interactions rather than all of them). This is an open research problem that future work would need to solve to fully verify the DSB's monotonicity claim. The paper's framework is fundamentally constrained by the curse of dimensionality in moment tensors, and this constraint will limit how thoroughly the DSB can be tested on high-dimensional real-world data regardless of methodology.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper reorients how we think about neural network learning trajectories. Rather than treating training as a monolithic process that produces a final function, it demonstrates that the temporal order in which networks extract statistical structure from data is predictable and universal—and that this order can be diagnosed with relatively simple synthetic probes that the network never sees during training. This is not a paradigm shift in the sense of overturning an established theory, but it is a significant methodological and conceptual reframing of how to study generalization dynamics. The paper gives the field a standardized diagnostic toolkit—optimal transport for grafting, maximum-entropy sampling for deletion—that can be applied to any trained model, any dataset, any modality, without modifying the training procedure.

Reconciling contradictory findings. Prior work on simplicity biases was fragmented across multiple explanatory frameworks that often seemed disjoint. The Frequency Principle (Xu et al., 2019b) described learning order in the spectral domain. The Neural Tangent Kernel literature (Jacot et al., 2018; Canatar et al., 2021) characterized the inductive biases of infinite-width networks through kernel eigenfunctions. Architectural simplicity bias (Valle-Perez et al., 2018; Chiang et al., 2023) attributed generalization to the parameter-function map rather than the optimizer. The DSB in this paper provides a unifying decomposition that sits alongside these frameworks: it describes the same phenomenon (networks learn simple structure first) but through the lens of distributional statistics rather than frequency, kernel eigenfunctions, or parameter-space priors. The key move is the Taylor expansion of expected loss (Equation 2), which connects moments of the data distribution directly to the loss surface that gradient descent navigates. This doesn't replace the Frequency Principle or NTK theory—it complements them by offering a decomposition that is particularly natural for understanding distribution-level sensitivity and that extends cleanly across modalities.

The paper also implicitly reconciles the tension between architectural simplicity bias (the claim that architectures inherently favor simple functions, regardless of optimizer) and optimization-dependent simplicity bias (the claim that gradient descent finds simple functions first). The finding that the DSB's U-shaped curves are nearly identical across ConvNeXt V2, Swin Transformer V2, and RegNet-Y—three very different architectures—and across SGD with momentum and AdamW, suggests that the temporal ordering is not architecture-specific (ruling out a pure architectural explanation) but also not optimizer-specific (ruling out a narrow optimization-dynamics explanation). Instead, the DSB appears to be a robust emergent property of gradient-based training on structured data, likely arising from the interplay between the factorial decay of Taylor coefficients at initialization and the greedy local improvement of gradient descent. If both the architectural simplicity bias and the DSB are correct, they operate at different levels: architecture constrains the space of learnable functions, while the DSB determines the order in which learnable functions are discovered within that space.

Research directions that become more attractive. The most immediate consequence of this paper is that probing for statistical sensitivity becomes a standard diagnostic, much like evaluating on adversarial examples or measuring calibration. The paper's optimal transport and maximum-entropy methods are computationally tractable (Appendix G: ~65 seconds on an L40 GPU for CIFARNet second-order probes) and the code is released, lowering the barrier to adoption. Future work on any new architecture, training recipe, or dataset can—and arguably should—include a DSB probe suite as part of standard evaluation. This makes questions like "does this new normalization scheme change the order in which statistics are learned?" or "does this self-supervised objective accelerate or decelerate the progression through statistical orders?" directly answerable with a few GPU-hours.

The double-descent finding in language models also makes in-context learning research more tightly connected to the simplicity bias literature. Before this work, the two phenomena were studied largely independently—simplicity bias as a property of parametric learning during training, in-context learning as a capability that emerges at scale. Figure 7 shows they interact: the DSB pushes loss up on n-gram sequences as the model moves beyond low-order statistics, while in-context learning pushes loss down as the model infers the n-gram data-generating process on the fly. This opens a research program on how parametric and non-parametric learning modes trade off during training, and whether similar interactions appear for other simple data-generating processes beyond n-gram models.

Research directions that become less attractive. The paper's finding that "model scale has a remarkably small effect on the learning curves" for the DSB's U-shaped pattern (Section 3.3) suggests that simply scaling up models does not fundamentally alter the order in which they learn statistical structure—at least within the range tested (Atto to Tiny for vision, 14M to 12B for language). Researchers hoping that larger models would bypass the "low-order statistics first" phase and immediately exploit high-order structure will be disappointed by this data. The DSB appears to be a robust property of the learning dynamics rather than a capacity constraint. This also means that scaling laws research, which typically focuses on final performance as a function of compute, may need to incorporate temporal dimensions—knowing that a 10× larger model achieves better final accuracy doesn't tell you whether it spends proportionally more or less of its training trajectory in the "low-order statistics" regime. The DSB provides a lens for studying this question.

Follow-Up Research This Work Enables

Mechanistic validation of the Taylor expansion argument. The paper's Equation 2 provides a theoretical motivation for the DSB but is never empirically validated. A strong follow-up would directly measure the evolution of the terms in this expansion during training. Specifically: for a trained network at multiple checkpoints, compute the partial derivatives $(\partial^\alpha L)(\mu)$ at the mean input $\mu$ for all multi-indices up to, say, order 4, and the corresponding central moments $\mathbb{E}[(x - \mu)^\alpha]$ of the data distribution. Then compute the contribution of each order $k$ to the expected loss—the sum over all $|\alpha| = k$ of $\frac{(\partial^\alpha L)(\mu)}{\alpha!} \mathbb{E}[(x - \mu)^\alpha]$—and track how these contributions change across training steps. The DSB predicts that the order-1 contribution should be reduced earliest, followed by order-2, then order-3, etc. This would require computing high-dimensional derivatives (feasible with automatic differentiation for moderate $d$ and low $k$) and would directly test whether the factorial-coefficient argument holds for real networks. A negative result—finding that higher-order terms are reduced before lower-order terms—would seriously challenge the DSB's mechanistic basis even if the behavioral U-shaped curves persist.

Causal test of the in-context learning driver of double descent. The paper attributes the second descent in Pythia's n-gram loss curves to in-context learning based on correlational timing evidence (comparing Figures 6 and 7). A clean causal experiment would evaluate Pythia checkpoints on n-gram sequences where in-context learning is impossible and check whether the second descent disappears. Two designs: (1) evaluate on n-gram sequences of length 1 (single tokens), where in-context learning cannot operate because there is no context; the DSB alone should produce a simple U-shaped loss curve with no second descent; (2) evaluate on n-gram sequences where token positions are randomly shuffled—this preserves the unigram/bigram/trigram frequencies but breaks the Markov structure, so the data-generating process is no longer inferrable from context. If the second descent is caused by in-context learning, it should vanish for shuffled sequences. This experiment would cost almost nothing computationally (evaluating existing checkpoints on new data) and would directly confirm or refute the paper's central interpretation of its most novel finding.

Skip-gram-inclusive probing of language models. The paper's Theorems 2.1-2.2 are stated for the full definition of n-gram statistics including skip-grams, but experiments use only contiguous n-grams. A natural extension would construct maximum-entropy models that match a richer subset of n-gram statistics—not all skip-grams (which is combinatorially infeasible), but a tractable subset that includes both contiguous and non-contiguous token co-occurrences at specific lag distances. For example, a "lag-k" bigram model conditions each token on the token at exactly $k$ positions prior, for multiple values of $k$. Comparing Pythia's loss curves on contiguous-only vs. lag-inclusive probes would reveal whether the network learns long-range token co-occurrences on a different schedule than adjacent bigrams. If the DSB holds, lagged bigram loss should follow a similar temporal pattern to contiguous bigram loss—both being second-order moments—but potentially with a later peak if long-range correlations are harder to learn. This experiment would bridge the gap between the theory (which includes skip-grams) and the current experiments (which exclude them), and the computational cost is moderate: a lag-1 bigram model has exactly the same size transition table as a standard bigram model, just with the conditioning token offset by 1 position.

Per-example difficulty analysis of the DSB U-shape. The paper reports only aggregate accuracy/loss curves averaged over all test images. A natural extension—directly parallel to the reference paper's per-difficulty-bin analysis—would bin images by how "gatherable" they are at each statistical order and analyze whether the U-shape is uniform or compositional. Specifically: for each test image, compute the change in the network's predicted class probability when that image is edited via first-order CQN or second-order Gaussian OT to match a different class. Images where this change is large are "highly gatherable" at that order; images where it is small are "resistant." Then bin images by their gatherability score and plot the U-shaped accuracy curves separately per bin. If the DSB is uniform, all bins should show similar U-shapes (just at different absolute accuracy levels). If the DSB is compositional—driven by a subset of images that happen to have class-distinctive means and covariances—the high-gatherability bins will show a strong U-shape while the low-gatherability bins show none. This experiment would provide the first evidence about whether the DSB describes a universal property of all inputs or a statistical average over heterogeneous behavior. It requires only re-analyzing the existing probe data with per-image metrics, not generating new synthetic data.

Extension to self-supervised and contrastive learning paradigms. All experiments in this paper use supervised classification or autoregressive language modeling. Whether the DSB holds under self-supervised objectives (SimCLR, MoCo, MAE) or contrastive language-image pretraining (CLIP) is completely open. The probing methodology transfers directly: train a SimCLR model on CIFAR-10, freeze the encoder at log-spaced checkpoints, train a linear probe on the real training set to establish classification performance, then evaluate the frozen encoder's representations on maximum-entropy and OT-edited images using the same linear probe. The DSB would predict that early SimCLR checkpoints have representations that are more sensitive to low-order statistics (the linear probe classifies OT-edited images by their grafted statistics better than chance) while later checkpoints lose this sensitivity. A negative result—finding no U-shape for self-supervised models—would suggest that the DSB is specific to supervised objectives that directly optimize for class separation, rather than a general property of representation learning. This would be important for understanding the scope of the DSB.

Deliberate counteraction of the DSB through adversarial moment training. If the DSB is a soft bias rather than a hard constraint, it should be possible to counteract it by training with a loss that explicitly penalizes reliance on low-order moments. A concrete experiment: during training, for each minibatch, generate CQN-edited and Gaussian OT-edited versions of the images (with random target classes), and add an auxiliary loss term that penalizes the network for classifying these edited images according to their grafted statistics. If the DSB can be overridden, the U-shaped accuracy curve on the probe data should be substantially reduced or eliminated in the adversarially trained model compared to a standard-trained control. This experiment would also have practical implications: if adversarial moment training produces models that are less sensitive to distribution shifts that alter low-order statistics (and instead rely on more robust higher-order features), it could serve as a data augmentation technique for improving out-of-distribution generalization. The computational cost is moderate—the OT maps are precomputed per class pair, and editing images during training adds only a matrix multiplication per image.

Practical Applications and Downstream Use Cases

Diagnosing model robustness to distribution shift. A model that relies primarily on low-order statistics (means and covariances) will be fragile under distribution shifts that alter those statistics—e.g., a change in lighting conditions, a shift in the color palette of images, or a domain transfer from photos to sketches. Conversely, such a model will be robust to shifts that preserve means and covariances but alter higher-order structure—e.g., applying a texture randomization filter. The paper's optimal transport and maximum-entropy probes can be used at any point in training to assess which statistical orders a given model checkpoint is sensitive to, without requiring access to the target distribution. For a practitioner deploying a model in a new environment where the exact nature of the distribution shift is unknown, running the DSB probe suite on the deployed checkpoint would reveal which statistical interventions change the model's predictions, flagging potential vulnerabilities. For example, if a ConvNeXt V2 model at deployment achieves 85% accuracy on CIFAR-10 but classifies 40% of Gaussian-OT-edited images as the target class (as in Figure 3 at step 2^10), the model is substantially relying on means and covariances and will likely fail under color/lighting shifts. This diagnostic costs nothing beyond generating the synthetic probe data once and evaluating the model, making it a lightweight addition to standard model evaluation pipelines.

Choosing early stopping points for transfer learning. The paper's finding that the DSB's U-shape peaks early in training—around step 2^9 to 2^10 for second-order Gaussian OT on CIFAR-10 (Figure 3)—suggests that early checkpoints are most sensitive to low-order statistics. For transfer learning to a target domain that differs from the source domain primarily in higher-order statistics (e.g., fine-tuning an ImageNet-pretrained model on medical images where textures differ but color distributions are similar), an early checkpoint might transfer better than a fully converged one, because the early checkpoint's reliance on low-order statistics aligns with the target domain's shared low-order structure. The paper's probe methodology provides a principled way to select the checkpoint: evaluate candidate checkpoints on maximum-entropy or OT-edited versions of the source-domain data, and select the checkpoint where accuracy is highest at the statistical order that is expected to transfer (e.g., second-order for color-invariant transfer). This is more systematic than the common practice of selecting the checkpoint with the best validation accuracy on the source domain, which may correspond to a model that has moved beyond transferable low-order statistics into source-specific higher-order features.

Data augmentation design informed by statistical sensitivity. The paper's dataset-dependence analysis (Section 3.3, Figure 2) reveals that CIFAR-10 requires higher-order statistics for good performance (second-order maximum-entropy samples look like colored noise), while MNIST can be nearly solved with second-order statistics alone (second-order samples are clearly recognizable digits). This suggests that the effectiveness of data augmentation techniques that randomize lower-order statistics (e.g., color jitter, Gaussian blur) vs. higher-order statistics (e.g., CutOut, random texture replacement) depends on the dataset. For CIFAR-10, aggressive color jitter early in training might help by forcing the network to move beyond first-order color cues, potentially accelerating the U-shaped progression and improving final robustness. For MNIST, the same color jitter would likely hurt, because color is not a discriminative cue to begin with and the network already relies primarily on second-order shape information. The DSB probe suite provides a diagnostic: if a dataset shows a deep U-shape on first-order probes (like CIFAR-10), first-order-randomizing augmentation is indicated; if first-order probes already plateau near validation accuracy (like MNIST), such augmentation is likely harmful. This could inform automated augmentation selection: run the DSB probes once on a new dataset, inspect the U-shape depth at each statistical order, and select augmentation strategies that randomize the orders where the network over-relies.