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∝m, while maintaining performance under adversarial shift requires additional in-context examples with Nρ−N0∝ρ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 Bρ(Q0) . The nominal distribution Q0 (blue) and an adversarial distribution Q (red) are shown. The black arrow represents the Wasserstein distance ρ 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 ρ and model capacity m . The surface illustrates how risk increases quadratically with ρ 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 ρmax with standard-error bars versus m . Middle: log-log plot showing an approximately linear power-law relationship. Right: collapse plot of ρmax/m versus m , which remains approximately flat across capacities. Despite anisotropy, heavy-tailed noise, and finite-sample effects, the scaling remains close to ρmax∝m .
Figure 4: Experiment 2: adversarial sample complexity across capacities. Left: ΔN=Nρ−N0 versus ρ2 for m∈{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 ΔN≈a(m)ρ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 ρ and the corresponding linear trend against ρ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 meff=mhead×L (head dimension times number of layers). The vertical axis reports the empirical safe radius αmax , the largest label-flipping fraction for which accuracy remains above 70%; this is the real-LLM analogue of ρ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 αmax≈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 m2 ), 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 ρ increases. Right: threshold-derived sample tax ΔN=Nρ−N0 plotted against ρ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 ρmax across various depths ( L∈{1,2,3,4} ) and capacity allocations. The fitted trend suggests that robustness scales with the square root of the accumulated capacity meff in these deep linear-Transformer experiments.
Figure 10: Effect of pretraining task diversity on empirical risk and structural robustness. Left: Broader pretraining (higher σβ ) 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 ρ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
ρ
Min m
Min N
Δ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 ρ (Wasserstein ball radius), the tool recommends minimum model capacity m and minimum in-context examples N . ΔN=Nρ−N0 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 ρ .
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 LN versus λN/N for four sample sizes, compared with the asymptotic formula 1/(1+λN/N) . Right: tightness ratio LNemp/LNtheory showing the bound is conservative but improves with larger N .
Model
m
Layers
ρmax ( N=4 )
ρmax ( N=8 )
Clean Acc. ( N=8 )
0.5B
64
24
0.37
0.44 ⋆
0.84
1.5B
128
28
0.40
0.51
0.97
3B
128
36
0.60 †
0.59
1.00
7B
128
28
0.57 †
0.47
0.98
Appendix
Table 3: Safe radius ρmax on Safety task for Qwen2.5-Instruct (mean over 3 seeds). ⋆ = head-dimension transition ( m jumps from 64 to 128). † = 3B > 7B despite fewer params (same m , 3B has more layers).
Model
ρmax (Base)
ρmax (Instruct)
Clean Acc. (Base)
Clean Acc. (Instruct)
0.5B
0.00 ▲
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 ρmax on Safety—Base vs Instruct ( N=4 ). ▲ = ICL threshold: Base 0.00 vs Instruct 0.37. Above threshold, Base ≈ Instruct.
Model Family
Model
m
ρ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 ρmax ( N=4 ) and clean accuracy for ICL-capable models (clean acc. ≥70% ). Models below threshold (all BLOOM, all OPT, small Pythia/Cerebras) show near-zero robustness regardless of capacity.
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.
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
Adversarial training (AT) is an effective defense for large language models (LLMs) against jailbreak attacks, but performing AT on LLMs is costly. To improve the efficiency of AT for LLMs, recent studies propose continuous AT (CAT) that searches for adversarial inputs within the continuous embedding space of LLMs during AT. While CAT has achieved empirical success, its underlying mechanism, i.e., why adversarial perturbations in the embedding space can help LLMs defend against jailbreak prompts synthesized in the input token space, remains unknown. This paper presents the first theoretical analysis of CAT on LLMs based on in-context learning (ICL) theory. For linear transformers trained with adversarial examples from the embedding space on in-context linear regression tasks, we prove a robust generalization bound that has a negative correlation with the perturbation radius in the embedding space. This clearly explains why CAT can defend against jailbreak prompts from the LLM's token space. Further, the robust bound shows that the robustness of an adversarially trained LLM is closely related to the singular values of its embedding matrix. Based on this, we propose to improve LLM CAT by introducing an additional regularization term, which depends on singular values of the LLM's embedding matrix, into the objective function of CAT. Experiments on real-world LLMs demonstrate that our method can help LLMs achieve a better jailbreak robustness-utility tradeoff. The code is available at https://github.com/fshp971/continuous-adv-icl.
Shaopeng Fu, Di Wang
Provable Responsible AI and Data Analytics (PRADA) Lab, King Abdullah University of Science and Technology, Saudi Arabia
Adversarial training is one of the most reliable defenses against adversarial attacks, but its high computational cost must generally be paid anew for each task. Robust foundation models offer a promising alternative: adversarially pretrain a model once and then transfer its robustness to downstream tasks through lightweight adaptation. However, a fundamental question remains open: can robustness acquired during pretraining transfer to unseen tasks without further adversarial training? In this study, we answer this question affirmatively. A single model adversarially pretrained at scale can achieve optimal robustness on new tasks without additional task-specific training. Specifically, we show that, for a family of Gaussian-mixture classification tasks, a sufficiently deep linear transformer adversarially trained across tasks can asymptotically attain the robust Bayes error on previously unseen tasks through in-context learning from clean demonstrations. By contrast, a standardly trained model cannot. We further analyze convergence under gradient flow, an accuracy--robustness trade-off, and demonstration complexity.