Bigger Is Safer: Provable Robustness in In-Context Learning Scales with Capacity

Di Zhang; Ningxu Zhang; Zimeng Liu

cs.LGFeb 19, 2026

Bigger Is Safer: Provable Robustness in In-Context Learning Scales with Capacity

Authors: Di Zhang, Ningxu Zhang, Zimeng Liu

Organizations: School of AI and Advanced Computing Xi’an Jiaotong-Liverpool University

Abstract

In-context learning (ICL) allows large language models to adapt to new tasks from a few examples without updating their parameters. Existing theories explain ICL by assuming the test task distribution matches pretraining -- an assumption that breaks down under adversarial distribution shifts. We introduce a distributionally robust meta-learning framework that provides worst-case guarantees for ICL under Wasserstein-based distribution shifts. Focusing on linear self-attention Transformers, we derive a non-asymptotic bound connecting adversarial perturbation strength ( $ρ$ ), model capacity ( $m$ ), and the number of in-context examples ( $N$ ). The analysis reveals that the maximum safe perturbation radius scales as $ρ_{\max} \propto \sqrt{m}$ , while maintaining performance under adversarial shift requires additional in-context examples with $N_ρ- N_0 \propto ρ^2$ . Experiments on synthetic tasks confirm these scaling laws, and experiments on 21 real pretrained models (0.1B--7B parameters, 5 families) provide qualitative evidence consistent with the theory's predictions, while revealing that ICL capability is a prerequisite for robustness. These findings advance the theoretical understanding of ICL under adversarial conditions and formalize the sense in which larger models are safer under distributional shift.

Figures & tables

Figure 1: Conceptual illustration of an adversarial distribution shift within the Wasserstein ball $\mathcal{B}_{\rho}(\mathbb{Q}_{0})$ . The nominal distribution $\mathbb{Q}_{0}$ (blue) and an adversarial distribution $\mathbb{Q}$ (red) are shown. The black arrow represents the Wasserstein distance $\rho$ between them. The background schematic connects task points to a Transformer via attention lines.

Figure 2: Schematic visualization of the worst-case meta-risk upper bound as a function of adversarial radius $\rho$ and model capacity $m$ . The surface illustrates how risk increases quadratically with $\rho$ but becomes progressively flatter as $m$ increases, demonstrating the mitigating effect of model capacity on adversarial vulnerability.

Figure 3: Experiment 1: robustness-capacity scaling under non-ideal conditions. Left: mean $\rho_{\max}$ with standard-error bars versus $\sqrt{m}$ . Middle: log-log plot showing an approximately linear power-law relationship. Right: collapse plot of $\rho_{\max}/\sqrt{m}$ versus $m$ , which remains approximately flat across capacities. Despite anisotropy, heavy-tailed noise, and finite-sample effects, the scaling remains close to $\rho_{\max}\propto\sqrt{m}$ .

Figure 4: Experiment 2: adversarial sample complexity across capacities. Left: $\Delta N=N_{\rho}-N_{0}$ versus $\rho^{2}$ for $m\in\{16,32,64,128,256\}$ in a shared coordinate system, with standard-error bars across seeds. The curves remain qualitatively parallel but are no longer identical. Right: fitted quadratic coefficient $a(m)$ from $\Delta N\approx a(m)\rho^{2}+b(m)$ as a function of model capacity, now evaluated over five capacities to make the trend more visible. The coefficient decreases systematically with model capacity, indicating a weaker adversarial sample tax for larger models.

Figure 5: Model design tradeoffs under budgeted allocation and safety-constrained deployment. Figure 5a–5b (Scenario A): the allocation-tradeoff main view and the utility-aware frontier, showing how nominal risk, adversarial risk, and compute cost interact under a fixed budget. Figure 5c–5d (Scenario B): the minimum required capacity as a function of threat level $\rho$ and the corresponding linear trend against $\rho^{2}$ , yielding a direct deployment-sizing rule.

Figure 6: Cross-family validation of the robustness-capacity relation on 21 pretrained language models. The horizontal axis is effective capacity

m_{\mathrm{eff}}=m_{\mathrm{head}}\times L

(head dimension times number of layers). The vertical axis reports the empirical safe radius

\alpha_{\max}

, the largest label-flipping fraction for which accuracy remains above 70%; this is the real-LLM analogue of

\rho_{\max}

in the theory. Marker color encodes head dimension, marker shape encodes model family, and marker size encodes clean accuracy. Models that fail the ICL prerequisite concentrate near

\alpha_{\max}\approx 0

, while ICL-capable models show a positive association between effective capacity and safe radius.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Module	Paper Theorem(s)
Basic.lean	Foundational real-analysis helpers
RobustnessBound.lean	Lemma 3.1 (ridge equivalence),
	Lemma 3.3 (Lipschitz bound)
MainTheorems.lean	Theorem 3.4 (worst-case meta-risk)
Corollaries.lean	Corollaries 3.5 , 3.6 ,
	and related algebraic implications

Appendix

Table 1: Mapping between Lean modules and paper theorems.

Figure 7: Efficiency of defense strategies. Left: Pareto frontier of Nominal vs Adversarial Risk. Traditional defenses suffer from a stronger trade-off, whereas capacity expansion reduces adversarial risk with a milder nominal-risk penalty in our experiments. Right: Adversarial risk reduction per unit of training compute. This panel reflects training-FLOPs efficiency only; capacity expansion incurs higher inference cost (approximately scaling as $m^{2}$ ), so the comparison should not be interpreted as end-to-end compute efficiency.

Figure 8: Qualitative threshold diagnostics for the BERT-base SST-2 stress test. Left: continuous accuracy-versus-perturbation curves for multiple context sizes $N$ , showing smooth degradation as $\rho$ increases. Right: threshold-derived sample tax $\Delta N=N_{\rho}-N_{0}$ plotted against $\rho^{2}$ , together with a least-squares fit and confidence band. The step-like pattern arises from the discrete definition of the minimal sample size under a fixed 80% accuracy threshold, rather than contradicting the underlying quadratic scaling trend. This appendix experiment should therefore be interpreted as a qualitative stress test rather than a direct verification of the linear-theory assumptions.

Figure 9: Maximum tolerable perturbation radius $\rho_{\max}$ across various depths ( $L\in\{1,2,3,4\}$ ) and capacity allocations. The fitted trend suggests that robustness scales with the square root of the accumulated capacity $\sqrt{m_{eff}}$ in these deep linear-Transformer experiments.

Figure 10: Effect of pretraining task diversity on empirical risk and structural robustness. Left: Broader pretraining (higher $\sigma_{\beta}$ ) exponentially decreases the model’s nominal bias on downstream tasks. Right: Because high-diversity models start with vastly superior nominal performance, they enjoy a much wider safety buffer before hitting the maximum tolerable risk threshold, leading to a much larger safe radius $\rho_{\max}$ .

Figure 11: Deployment framework recommendations across four threat scenarios. Left: Required capacity and samples for each scenario. Right: Capacity vs. threat level with scenario labels and 90% confidence bands.

Scenario	$\rho$	Min $m$	Min $N$	$\Delta N$
Safe	0.30	3	28	5
Standard	0.60	12	41	18
Critical	1.00	34	74	51
Extreme	1.50	77	138	115

Appendix

Table 2: Deployment Framework Recommendations. For different threat levels

\rho

(Wasserstein ball radius), the tool recommends minimum model capacity

m

and minimum in-context examples

N

\Delta N=N_{\rho}-N_{0}

is the additional samples needed beyond the nominal setting. These are illustrative projections only.

Figure 12: Direct verification of Theorem 3.4 . Empirical worst-case risk (circles) versus the theoretical upper bound (dashed) for three capacities. The bound holds for all 189 test points (100% validity) with a mean relative gap of 36–74%, confirming that the bound is valid and reasonably tight for moderate $\rho$ .

Figure 13: Compression lower bound verification. Left: empirical error versus compression fraction, compared with the Fano lower bound. Right: error versus retained mutual information. Empirical error exceeds the Fano bound at all 20 points and increases monotonically with compression, approaching random guessing ( $1-1/K=0.875$ ) at high compression.

Figure 14: Lemma 3.3 verification. Left: empirical $L_{N}$ versus $\lambda_{N}/N$ for four sample sizes, compared with the asymptotic formula $1/(1+\lambda_{N}/N)$ . Right: tightness ratio $L_{N}^{\text{emp}}/L_{N}^{\text{theory}}$ showing the bound is conservative but improves with larger $N$ .

Model	$m$	Layers	$\rho_{\max}$ ( $N{=}4$ )	$\rho_{\max}$ ( $N{=}8$ )	Clean Acc. ( $N{=}8$ )
0.5B	64	24	0.37	0.44 $\star$	0.84
1.5B	128	28	0.40	0.51	0.97
3B	128	36	0.60 $\dagger$	0.59	1.00
7B	128	28	0.57 $\dagger$	0.47	0.98

Appendix

Table 3: Safe radius

\rho_{\max}

on Safety task for Qwen2.5-Instruct (mean over 3 seeds).

\star

= head-dimension transition (

m

jumps from 64 to 128).

\dagger

= 3B

>

7B despite fewer params (same

m

, 3B has more layers).

Model	$\rho_{\max}$ (Base)	$\rho_{\max}$ (Instruct)	Clean Acc. (Base)	Clean Acc. (Instruct)
0.5B	0.00 $\blacktriangle$	0.37	0.57	0.84
1.5B	0.53	0.40	1.00	0.97
3B	0.60	0.60	1.00	1.00

Appendix

Table 4: Safe radius

\rho_{\max}

on Safety—Base vs Instruct (

N{=}4

\blacktriangle

= ICL threshold: Base 0.00 vs Instruct 0.37. Above threshold, Base

\approx

Instruct.

Model Family	Model	$m$	$\rho_{\max}$	Clean Acc.
Qwen2.5-Inst.	0.5B	64	0.37	0.84
Qwen2.5-Inst.	1.5B	128	0.40	0.97
Qwen2.5-Inst.	3B	128	0.60	1.00
Qwen2.5-Inst.	7B	128	0.57	0.98
Cerebras-GPT	111M	64	0.13	0.74
Cerebras-GPT	256M	64	0.27	0.85

Appendix

Table 5: Cross-family validation: safe radius

\rho_{\max}

(

N{=}4

) and clean accuracy for ICL-capable models (clean acc.

\geq 70\%

). Models below threshold (all BLOOM, all OPT, small Pythia/Cerebras) show near-zero robustness regardless of capacity.

In-Context Learning

In-context learning studies how models adapt to new tasks from examples or instructions supplied at inference time, without changing their weights.

See all 404 papers

How Does the Pretraining Distribution Shape In-Context Learning? A Fundamental Trade-Off
Oct 1, 2025Waïss Azizian, Ali HasanLanguage Model Pretraining OOD Generalization
Univ. Grenoble Alpes, CNRS, Inria, Grenoble INP, LJK, 38000 Grenoble, France · Work done during an internship at Morgan Stanley Machine Learning Research. · Machine Learning Research, Morgan Stanley, New York, USA
Understanding and Improving Continuous Adversarial Training for LLMs via In-context Learning Theory
Apr 14, 2026Shaopeng Fu, Di WangAdversarial Training Adversarial Robustness of Language Models
Provable Responsible AI and Data Analytics (PRADA) Lab, King Abdullah University of Science and Technology, Saudi Arabia
Adversarially Trained Linear Transformers Are Optimal Robust In-Context Learners for Gaussian Mixtures
Oct 6, 2026Soichiro KumanoAdversarial Training Transformer
LY Corporation

Oct 1, 2025J/K move · Enter open · S save

How Does the Pretraining Distribution Shape In-Context Learning? A Fundamental Trade-Off

Waïss Azizian, Ali Hasan

Univ. Grenoble Alpes, CNRS, Inria, Grenoble INP, LJK, 38000 Grenoble, France · Work done during an internship at Morgan Stanley Machine Learning Research. · Machine Learning Research, Morgan Stanley, New York, USA

Language Model Pretraining OOD Generalization In-Context Learning Heavy-Tailed Distributions

Open paper AskPDF arXiv

The factors driving the performance of in-context learning (ICL) in large language models (LLMs) remain poorly understood despite ICL's surprising effectiveness, enabling models to adapt to new tasks from only a handful of examples. To clarify and improve these capabilities, we characterize how the statistical properties of the pretraining distribution (e.g., tail behavior, coverage) shape ICL. We develop a theoretical framework that encompasses generalization and task selection and show how distributional properties govern sample efficiency, task retrieval, and robustness. To this end, we generalize existing concentration results to heavy-tailed priors and dependent sequences, better reflecting the structure of LLM pretraining data. Our framework reveals a fundamental design trade-off: heavy-tailed pretraining distributions facilitate robust task selection under distribution shifts but are detrimental to generalization, especially in low-data regimes. We then empirically evaluate our predictions by studying how ICL performance varies with the pretraining distribution on challenging tasks such as stochastic differential equations and stochastic processes with memory. Together, these findings suggest that controlling key statistical properties of the pretraining distribution is essential for building ICL-capable and reliable LLMs.

Abstract

In-Context Learning

Explore similar work

How Does the Pretraining Distribution Shape In-Context Learning? A Fundamental Trade-Off

Understanding and Improving Continuous Adversarial Training for LLMs via In-context Learning Theory

Adversarially Trained Linear Transformers Are Optimal Robust In-Context Learners for Gaussian Mixtures