arXiv subjects

Yujin Song

Publications and source records attributed to Yujin Song.

7 recordsLinked to original sources

arXiv2025

From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers

Transformers can implement both generalizable algorithms (e.g., induction heads) and simple positional shortcuts (e.g., memorizing fixed output positions). In this work, we study how the choice of pretraining data distribution steers a shallow transformer toward one behavior or the other. Focusing on a minimal trigger-output prediction task -- copying the token immediately following a special trigger upon its second occurrence -- we present a rigorous analysis of gradient-based training of a single-layer transformer. In both the infinite and finite sample regimes, we prove a transition in the learned mechanism: if input sequences exhibit sufficient diversity, measured by a low ``max-sum'' ratio of trigger-to-trigger distances, the trained model implements an induction head and generalizes to unseen contexts; by contrast, when this ratio is large, the model resorts to a positional shortcut and fails to generalize out-of-distribution (OOD). We also reveal a trade-off between the pretraining context length and OOD generalization, and derive the optimal pretraining distribution that minimizes computational cost per sample. Finally, we validate our theoretical predictions with controlled synthetic experiments, demonstrating that broadening context distributions robustly induces induction heads and enables OOD generalization. Our results shed light on the algorithmic biases of pretrained transformers and offer conceptual guidelines for data-driven control of their learned behaviors.

cs.LG↗

arXiv2025

Sample-efficient quantum error mitigation via classical learning surrogates

The pursuit of practical quantum utility on near-term quantum processors is critically challenged by their inherent noise. Quantum error mitigation (QEM) techniques are leading solutions to improve computation fidelity with relatively low qubit-overhead, while full-scale quantum error correction remains a distant goal. However, QEM techniques incur substantial measurement overheads, especially when applied to families of quantum circuits parameterized by classical inputs. Focusing on zero-noise extrapolation (ZNE), a widely adopted QEM technique, here we devise the surrogate-enabled ZNE (S-ZNE), which leverages classical learning surrogates to perform ZNE entirely on the classical side. Unlike conventional ZNE, whose measurement cost scales linearly with the number of circuits, S-ZNE requires only constant measurement overhead for an entire family of quantum circuits, offering superior scalability. Theoretical analysis indicates that S-ZNE achieves accuracy comparable to conventional ZNE in many practical scenarios, and numerical experiments on up to 100-qubit ground-state energy and quantum metrology tasks confirm its effectiveness. Our approach provides a template that can be effectively extended to other quantum error mitigation protocols, opening a promising path toward scalable error mitigation.

quant-ph↗

arXiv2025

How Does Label Noise Gradient Descent Improve Generalization in the Low SNR Regime?

The capacity of deep learning models is often large enough to both learn the underlying statistical signal and overfit to noise in the training set. This noise memorization can be harmful especially for data with a low signal-to-noise ratio (SNR), leading to poor generalization. Inspired by prior observations that label noise provides implicit regularization that improves generalization, in this work, we investigate whether introducing label noise to the gradient updates can enhance the test performance of neural network (NN) in the low SNR regime. Specifically, we consider training a two-layer NN with a simple label noise gradient descent (GD) algorithm, in an idealized signal-noise data setting. We prove that adding label noise during training suppresses noise memorization, preventing it from dominating the learning process; consequently, label noise GD enjoys rapid signal growth while the overfitting remains controlled, thereby achieving good generalization despite the low SNR. In contrast, we also show that NN trained with standard GD tends to overfit to noise in the same low SNR setting and establish a non-vanishing lower bound on its test error, thus demonstrating the benefit of introducing label noise in gradient-based training.

cs.LG↗

arXiv2024

Pretrained transformer efficiently learns low-dimensional target functions in-context

Transformers can efficiently learn in-context from example demonstrations. Most existing theoretical analyses studied the in-context learning (ICL) ability of transformers for linear function classes, where it is typically shown that the minimizer of the pretraining loss implements one gradient descent step on the least squares objective. However, this simplified linear setting arguably does not demonstrate the statistical efficiency of ICL, since the pretrained transformer does not outperform directly solving linear regression on the test prompt. In this paper, we study ICL of a nonlinear function class via transformer with nonlinear MLP layer: given a class of \textit{single-index} target functions $f_*(\boldsymbol{x}) = \sigma_*(\langle\boldsymbol{x},\boldsymbol{\beta}\rangle)$, where the index features $\boldsymbol{\beta}\in\mathbb{R}^d$ are drawn from a $r$-dimensional subspace, we show that a nonlinear transformer optimized by gradient descent (with a pretraining sample complexity that depends on the \textit{information exponent} of the link functions $\sigma_*$) learns $f_*$ in-context with a prompt length that only depends on the dimension of the distribution of target functions $r$; in contrast, any algorithm that directly learns $f_*$ on test prompt yields a statistical complexity that scales with the ambient dimension $d$. Our result highlights the adaptivity of the pretrained transformer to low-dimensional structures of the function class, which enables sample-efficient ICL that outperforms estimators that only have access to the in-context data.

cs.LG↗

arXiv2024

Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge combinations

We study the computational and sample complexity of learning a target function $f_*:\mathbb{R}^d\to\mathbb{R}$ with additive structure, that is, $f_*(x) = \frac{1}{\sqrt{M}}\sum_{m=1}^M f_m(\langle x, v_m\rangle)$, where $f_1,f_2,...,f_M:\mathbb{R}\to\mathbb{R}$ are nonlinear link functions of single-index models (ridge functions) with diverse and near-orthogonal index features $\{v_m\}_{m=1}^M$, and the number of additive tasks $M$ grows with the dimensionality $M\asymp d^\gamma$ for $\gamma\ge 0$. This problem setting is motivated by the classical additive model literature, the recent representation learning theory of two-layer neural network, and large-scale pretraining where the model simultaneously acquires a large number of "skills" that are often localized in distinct parts of the trained network. We prove that a large subset of polynomial $f_*$ can be efficiently learned by gradient descent training of a two-layer neural network, with a polynomial statistical and computational complexity that depends on the number of tasks $M$ and the information exponent of $f_m$, despite the unknown link function and $M$ growing with the dimensionality. We complement this learnability guarantee with computational hardness result by establishing statistical query (SQ) lower bounds for both the correlational SQ and full SQ algorithms.

cs.LG↗

arXiv2024

Structural Preprocessing Method for Nonlinear Differential-Algebraic Equations Using Linear Symbolic Matrices

Differential-algebraic equations (DAEs) have been used in modeling various dynamical systems in science and engineering. Several preprocessing methods for DAEs, such as consistent initialization and index reduction, use structural information on DAEs. Unfortunately, these methods may fail when the system Jacobian, which is a functional matrix, derived from the DAE is singular. To transform a DAE with a singular system Jacobian into a nonsingular system, several regularization methods have been proposed. Most of all existing regularization methods rely on symbolic computation to eliminate the system Jacobian for finding a certificate of singularity, resulting in much computational time. Iwata--Oki--Takamatsu (2019) proposed a method (IOT-method) to find a certificate without symbolic computations. The IOT method approximates the system Jacobian by a simpler symbolic matrix, called a layered mixed matrix, which admits a fast combinatorial algorithm for singularity testing. However, it often overlooks the singularity of the system Jacobian since the approximation largely discards algebraic relationships among entries in the original system Jacobian. In this study, we propose a new regularization method extending the idea of the IOT method. Instead of layered mixed matrices, our method approximates the system Jacobian by more expressive symbolic matrices, called rank-1 coefficient mixed (1CM) matrices. This makes our method more widely applicable. We give a fast combinatorial algorithm for finding a singularity certificate of 1CM-matrices, which is free from symbolic elimination. Our method is also advantageous in that it globally preserves the solution set to the DAE. Through numerical experiments, we confirmed that our method runs fast for large-scale DAEs from real instances.

cs.SC↗

arXiv2023

Angularity in Higgs boson decays via $\boldsymbol{H\to gg}$ at NNLL$'$ accuracy

We present improved predictions of a class of event-shape distributions called angularity for a contribution from an effective operator $H\to gg$ in Higgs hadronic decay that suffers from large perturbative uncertainties. In the frame of Soft-Collinear Effective Theory, logarithmic terms of the distribution are resummed at NNLL$'$ accuracy, for which 2-loop constant of gluon-jet function for angularity is independently determined by a fit to fixed-order distribution at NLO corresponding to ${\mathcal{O}}(\alpha_s^2)$ relative to the Born rate. Our determination shows reasonable agreement with the value in a thesis recently released. In the fit, we use an asymptotic form with a fractional power conjectured from recoil corrections at one-loop order and it improves the accuracy of determination in positive values of angularity parameter $a$. The resummed distribution is matched to the NLO fixed-order results to make our predictions valid at all angularity values. We also discuss the first moment and subtracted moment of angularity as a function of $a$ that allow to extract information on leading and subleading nonperturbative corrections associated with gluons.

hep-ph↗