Trapped by simplicity: When Transformers fail to learn from noisy features

AI summary

Overview

• Research area: Machine learning theory — noise-robust learning of Boolean functions, transformer inductive biases (simplicity bias), and connections to information/noisy-channel coding.
• Technical level: Intermediate (the core empirical claims are accessible; the supporting arguments use Boolean Fourier analysis, sensitivity, and the noise operator).
• Scope: An empirical and theoretical study of whether transformers trained on data with feature noise can still recover a target Boolean function that predicts correctly on noiseless inputs, compared against LSTMs.

What This Paper Is About

Noise is pervasive in the data used to train large language models, but it is unclear whether models trained on noisy features can still learn a target function that makes correct predictions on clean, noiseless inputs — a property the authors call noise-robust learning. The paper studies this question in the simplified setting of Boolean functions of binary inputs, where an optimal predictor for the noisy distribution can be computed exactly and compared against the target function. The authors find that transformers are sometimes excellent at this task (sparse parities and odd majorities) but generally fail for random k-juntas, and they attribute the failure to transformers' simplicity bias combined with the fact that the optimal noisy-data predictor is usually less sensitive than the target function.

Key Contributions

1. The authors show that transformers succeed at noise-robust learning for sparse parity and majority functions at high rates of feature noise, while LSTMs generally fail even at low noise levels.
2. They show that transformers fail at noise-robust learning for random k-juntas, even while simultaneously reaching near-optimal accuracy on noisy validation data.
3. They propose and test an explanation: the sensitivity of the optimal solution for noise-robust learning is rarely greater than the sensitivity of the target function, so transformers' simplicity bias drives them to a solution that is suboptimal for noiseless evaluation.
4. They demonstrate a controlled "trap" function where transformers converge to an incorrect low-sensitivity solution, and show transformers can escape the trap when trained with an additional loss term that penalizes high-sensitivity solutions.

Main Findings

• Transformers win on parities and odd majorities: For maj(20,5) and maj(40,5) (with N = 2000 noisy training examples), the median transformer reliably outperformed the best LSTM across 300 training runs per model with varied hyperparameters and random initializations; validation accuracy was estimated using 10000 examples.
• Sparse parity is harder but still favors transformers: Both LSTMs and SANs fail in a large fraction of training experiments on parity(20,4) with feature noise, but transformers successfully learn it — defined as achieving noiseless accuracy ≥ 95% — more often than LSTMs, even when both architectures perform comparably at zero noise rate. The behavior also extended to sparse multitask parities, where transformers succeeded and all LSTMs failed.
• Why those functions are special: Proposition 1 states that for f ∈ {maj_n, parity} with odd n, the function itself is optimal for prediction on noisy features (f = f_N*); such functions are "self-predicting" in the terminology of Weinberger and Shayevitz (2018), and can therefore be learned by ordinary loss minimization. The authors call this property atypical.
• Random k-juntas fail: Across 3200 randomly generated k-juntas (with k ∈ (5,6,7,8)), transformers performed worse at noise-robust learning as the sensitivity gap (sens[f] − sens[f_N*]) grew — despite achieving near-optimal validation error with little overfitting.
• The optimal noisy predictor is typically simpler: Proposition 2 shows that for a function sampled uniformly at random on n bits, the optimal predictor's average sensitivity is approximately (n/π) · arccos(2p(1−p)/(p² + (1−p)²)), and that E_f[sens[f]] > E_f[sens[f_N*]] for p ∈ (0, 1/2]. Simulations found no violations of sens[f] ≥ sens[f_N*] for random k-juntas, though the authors note the inequality can only hold as an average-case statement.
• Two models, two different failure modes: Using a trap function with err_f(f_N*) ≈ err_f(f) but sens[f_N*] ≪ sens[f], the authors show transformers converge toward the trap function, while LSTMs fail by overfitting to the training data — a completely different reason.
• A sensitivity penalty can break the trap: Adding an explicit loss penalty λ · sens[f̂] (with λ = 1) allowed transformers to learn f from noisy data for a narrow range of λ; the authors note there is a clear optimum for λ but that it may not be practical to optimize.
• Even majorities are a lost cause: For f := maj(n,k) with even n, f_N* ≠ f in general because f is imbalanced. Transformers typically failed to learn maj(30,4) from noisy features (at p = 0.32) for many penalty parameters λ, while LSTMs became more capable of learning the lower-sensitivity f_N*; the performance gap between the median SAN and best LSTM for even-n majorities almost completely disappears.

Methodology in Plain English

The authors work with uniformly random bitstrings X, apply a Boolean function f to produce a label Y = f(X), then corrupt the input bits with independent, symmetric bitflips at rate p to produce noisy features Z = X ⊕ E. Models are trained only on (Z, Y) pairs — never on label noise — and evaluated on clean (X, Y) data.

They compare self-attention network transformers (SANs) against LSTMs. Because learning outcomes are sensitive to initialization and hyperparameters, they repeat each configuration 300 times with random initializations and randomly drawn hyperparameters, chosen from a grid initially tuned on noiseless data and slightly expanded for noisy conditions. Reporting uses the median transformer versus the best LSTM, which is a deliberately conservative comparison favoring the LSTM.

To connect success or failure to simplicity, they use sensitivity (the expected number of single-bit flips that change a function's output) as the simplicity measure. The optimal predictor for the noisy distribution, f_N*(x) := sign(T_{1−2p} f(x)), is computed via the noise operator, letting them compare f and f_N* on sensitivity and noisy generalization error. Their central inequality err_f(f̂) ≤ ϵ + p · sens[f̂] explains why low sensitivity helps at evaluation time but hurts when the data themselves are noisy. Finally, they design a "trap" function where the target and the optimal noisy predictor have nearly identical noisy error but very different sensitivity, isolating the effect of the simplicity bias, and then add a sensitivity penalty to the loss to test whether the bias can be counteracted.

Why This Matters

• Impact on research: The work suggests that comparisons based only on noisy validation error can be misleading — a model can appear near-optimal while having learned a fundamentally different, simpler function than the target. It reframes "noise robustness" (robustness to noise at evaluation time) as distinct from "noise-robust learning" (learning from noisy training data), and connects next-token prediction to noisy channel coding rather than to compression.
• Real-world applications (as described in the paper):
  • Decoding classical error-correcting codes.
  • Decoding quantum error-correcting codes.
  • Arithmetic and discrete mathematics reasoning by language models, where each next token depends sensitively on preceding noiseless text.
  • General binary-domain classification and generative modeling where feature noise is prevalent.
• Industry relevance: The findings imply that LLMs trained on noisy, high-stochasticity text (incorrect grammar or semantics) may struggle to learn precise algorithmic tasks, and that reducing loss to the next-token conditional entropy could even be detrimental for learning exact concepts. The authors argue for mitigating simplicity biases in large language models if algorithmic tasks are to be learned from noisy training data.

Future Directions

1. Testing whether the key inequality sens[f] ≥ sens[f_N*] holds for realistic, correlated, or non-identically-distributed noise models; the authors note there is no reason to expect it to hold for noise models with arbitrary correlations between bitflips, although the mechanism of T_ρ does not depend on identical errors.
2. Determining whether these effects extend to domains with more complex inputs, such as natural language processing and modular arithmetic, where high-entropy training text may make sensitive function learning impossible.
3. Running further experiments on how noise and stochasticity in training data affect LLM reasoning at evaluation time, complementing prior observations that evaluation-time noise degrades arithmetic performance.
4. Developing more sophisticated techniques — beyond a fixed sensitivity penalty — for noise-robust learning in the regime where the optimal noisy predictor f_N* is much more accurate than the target f.

Target Audience

Machine learning theory researchers and practitioners interested in transformer inductive biases, learning theory for Boolean functions, and the reliability of models trained on noisy data. It is also relevant to engineers applying transformers to error correction or reasoning tasks, and to readers already familiar with prior work on transformers' simplicity bias and low-sensitivity preferences — the paper assumes comfort with Boolean analysis, entropy bounds, and the noise operator, with much of that background deferred to Appendix A.

Authors’ abstract

Noise is ubiquitous in data used to train large language models, but it is not well understood whether these models are able to correctly generalize to inputs generated without noise. Here, we study noise-robust learning: are transformers trained on data with noisy features able to find a target function that correctly predicts labels for noiseless features? We show that transformers succeed at noise-robust learning for a selection of $k$-sparse parity and majority functions, compared to LSTMs which fail at this task for even modest feature noise. However, we find that transformers typically fail at noise-robust learning of random $k$-juntas, especially when the boolean sensitivity of the optimal solution is smaller than that of the target function. We argue that this failure is due to a combination of two factors: transformers' bias toward simpler functions, combined with an observation that the optimal function for noise-robust learning typically has lower sensitivity than the target function for random boolean functions. We test this hypothesis by exploiting transformers' simplicity bias to trap them in an incorrect solution, but show that transformers can escape this trap by training with an additional loss term penalizing high-sensitivity solutions. Overall, we find that transformers are particularly ineffective for learning boolean functions in the presence of feature noise.