We analyze approximate message passing (AMP) with an independent Gaussian initialization for noiseless phase retrieval in the proportional asymptotic regime. A random initialization has overlap of order $d^{-1/2}$ with the signal, and AMP requires a growing number of iterations to attain non-vanishing overlap. Thus, its precise behavior cannot be characterized by classical fixed-time state evolution. We prove a Gaussian decomposition of the AMP trajectory and control its error over the horizons required for recovery. The resulting analysis shows that random initialization attains the weak-recovery threshold $δ_{\rm weak}=1/2$. For $δ\in(δ_{\rm weak},δ_{\rm str})$, where $δ_{\rm str}\approx1.13$, the signal strength follows state evolution and approaches its stable finite fixed point uniformly for \(n^{1/3}/\operatorname{polylog}(n)\) iterations. For $δ>δ_{\rm str}$, AMP reaches any prescribed fixed recovery accuracy within $O_{δ,\varepsilon}(\log n)$ iterations. The majority of our analysis applies more generally to generalized AMP for single-index models.
Bum Jun Kim, Kohei Hayashi, Shunsuke Kamiya +3cs.LG cs.CL
Looped Transformers create a useful train- and test-time compute axis by reusing the same Transformer block over recurrent depth, increasing effective depth at a fixed parameter count. However, that shared block must then govern an entire trajectory of varying hidden states over trained and extrapolated depths. Furthermore, in additive-injection looped Transformers, an input-conditioned signal is reintroduced at every recurrent step, so applying the shared transition at an input-conditioned reference can still move the hidden state. In this paper, we propose Source-Centered State Evolution (SCSE), which is designed to reconcile input conditioning with reference-preserving shared recurrence. Specifically, SCSE retains input dependence through its learned anchor and initial deviation, allows nonzero deviations to drive recurrent computation while mapping zero deviation to zero, and guarantees exact anchor invariance through its zero-deviation mask. The designated anchor is thereby a one-step fixed point by construction. The zero-deviation forcing bias is the next deviation produced from the anchor itself and vanishes in SCSE, while nonzero deviations remain active and support state-dependent recurrent computation. Our theory shows that the zero-deviation forcing bias is a design degree of freedom whose task effect can be harmful, neutral, or beneficial; SCSE resolves this choice in favor of exact anchor invariance by setting the bias to zero. Across WikiText-2, WikiText-103, direct web-corpus pretraining, held-out web-text transfer, and LAMBADA completion, SCSE improves the controlled recurrent quality frontier. Ablation studies identify the learned anchor and the anchor-coordinate deviation recurrence as the primary contributors to the gain, and a trained-model case study grounds the anchor-response diagnostic in observed recurrent motion.
In modern parametric model training, full-batch gradient descent (and its variants) suffers due to progressively stronger biasing towards the exact realization of training data; this drives the systematic ``generalization gap'', where the train error becomes an unreliable proxy for test error. Existing approaches either argue this gap is benign through complex analysis or sacrifice data to a validation set. In contrast, we introduce decoupled descent (DD), a novel theory-based training algorithm that satisfies a train-test identity -- enforcing the train error to asymptotically track the test error for stylized Gaussian mixture models. Within this specific regime, leveraging approximate message passing theory, DD iteratively cancels the biases due to data reuse, rigorously demonstrating the feasibility of zero-cost validation and $100\%$ data utilization. Moreover, DD is governed by a low-dimensional state evolution recursion, rendering the dynamics of the algorithm transparent and tractable. We validate DD on XOR classification, yielding superior performance compared to GD; additionally, we implement noisy MNIST and non-linear probing of CIFAR-10, demonstrating that even when our stylized assumptions are relaxed, DD narrows the generalization gap compared to GD.