Support Tokens, Stability Margins, and a New Foundation for Robust LLMs

Deepak Agarwal; Dhyey Dharmendrakumar Mavani; Suyash Gupta; Karthik Sethuraman; Tejas Dharamsi

cs.LGFeb 25, 2026

Support Tokens, Stability Margins, and a New Foundation for Robust LLMs

Authors: Deepak Agarwal, Dhyey Dharmendrakumar Mavani, Suyash Gupta, Karthik Sethuraman, Tejas Dharamsi

Organizations: LinkedIn, Mountain View, CA

Abstract

Self-attention is usually described as a flexible, content-adaptive way to mix a token with information from its past. We reinterpret causal self-attention transformers, the backbone of modern foundation models, within a probabilistic framework, much as classical PCA is extended to probabilistic PCA. This reformulation reveals a key structural consequence of the underlying change of variables: a barrier constraint emerges on the parameters of self-attention. The resulting geometry exposes a degeneracy boundary where the attention-induced mapping becomes locally ill-conditioned, yielding a stability-margin interpretation analogous to the margin in support vector machines. This, in turn, naturally gives rise to the concept of support tokens. We further show that causal transformers define a consistent stochastic process over infinite token sequences, providing a rigorous probabilistic foundation for sequence modeling. Building on this view, we derive a Bayesian MAP training objective that requires only a minimal modification to standard LLM training: adding a smooth log-barrier penalty to the usual cross-entropy loss. Empirically, the resulting training objective improves robustness to input perturbations and sharpens the margin geometry of the learned representations without sacrificing out-of-sample accuracy.

Figures & tables

Figure 1.1 : Margin-to-degeneracy geometry (SVM analogy). The curve denotes the degeneracy boundary: configurations where the attention-induced latent-noise mapping becomes locally ill-conditioned, so small perturbations produce unstable changes. A sequence traces a trajectory of token representations in embedding space across positions (dots). Schematically, a token’s margin is represented by its distance to this boundary (arrow). In this example, token 1 is the sequence’s support token because it lies closest to the boundary, so it controls the sequence-level stability margin and carries the largest barrier pressure. This mirrors the SVM picture, where support vectors are the points closest to the decision boundary and therefore determine the margin. A more precise formulation is given in Section 3 .

Figure 3.2 : Prior density under positive and negative coupling ( $d{=}1$ , $n{=}3$ ). (a) Conditional density of $x_{3}$ given $x_{1}{=}0$ , $x_{2}{=}2$ . Positive coupling ( $a{=}+0.25$ , red) broadens the baseline Gaussian distribution, while negative coupling ( $a{=}-0.25$ , green) sharpens it. (b) The diagonal change-of-variables factor $1-a\,\mathrm{Var}_{t}$ (Eq. 8 ) evaluated across $x_{3}$ . (c) Density profiles for varying coupling strengths $a\in\{-0.55,-0.35,0,+0.35,+0.55\}$ .

Figure 3.3 : Population-level behavior under positive vs. negative coupling ( $d{=}1$ , $n{=}5$ ). Evaluation of 4000 random sequences from $\mathcal{N}(0,4)$ under positive coupling ( $a{=}+0.2$ , top row) and negative coupling ( $a{=}-0.2$ , bottom row). (a, c) The diagonal change-of-variables factor as a function of the attention-weighted variance $\mathrm{Var}_{t}$ . Positive coupling introduces a hard degeneracy boundary at zero, whereas negative coupling remains strictly positive. (b, d) The resulting log-density. Under positive coupling, configurations near the degeneracy boundary are excluded (red points). Under negative coupling, the log-barrier is removed and larger dispersion is permitted.

Figure 7.4 : Training curves for cross-entropy-only vs. margin-regularized models. Perplexity over 20 epochs for the cross-entropy-only baseline (blue) and the margin-regularized model (green, $\lambda_{m}{=}0.05$ ). From left to right: wikitext-103 validation then simplebooks-92 validation. Across both datasets, the validation trajectories remain very close, indicating that the theory-derived margin penalty preserves predictive quality while adding a stability-oriented training signal.

Figure 7.5 : Robustness on wikitext-103 : relative degradation under unstructured Gaussian noise and structured Gaussian drift. (a) Unstructured Gaussian noise: grouped boxplots of perturbed perplexity / clean perplexity across 100 validation subsamples for $\sigma\in\{0,0.05,0.10,0.15,0.20,0.25\}$ . (b) Structured Gaussian drift: grouped boxplots for $\sigma\in\{0,1.5,3.0,4.5,6.0\}$ . In both panels, each box summarizes the empirical distribution across validation subsamples.

Figure 7.6 : Support-token concentration under the learned margin prior and a random-initialized baseline. For each dataset, we compute the exact per-position barrier score $b_{t}$ , convert it into a normalized distribution over positions, sort token positions by barrier weight, and plot the cumulative barrier pressure carried by the top- $k$ support tokens. Solid lines show the median across validation subsamples and shaded bands show the empirical 95% interval. Left: wikitext-103 . Right: simplebooks-92 . The dashed line marks $k{=}5$ .

Figure 7.7 : Selecting $\lambda_{m}$ on wikitext-103 . Grouped boxplots over 100 validation subsamples for $\lambda_{m}\in\{0.00,0.02,0.05,0.10,0.25\}$ . (a) Degradation ratio at fixed unstructured Gaussian noise $\sigma{=}0.25$ . (b) Contextual embedding error. (c) Effective support size.

Dataset	$\lambda_{m}$	Clean perplexity	Degradation ratio	Contextual embedding error	Effective support size
wikitext-103	0.00	15.68	775.81	0.608	—
wikitext-103	0.02	15.74	328.51	0.583	2.75
wikitext-103	0.05	15.71	364.92	0.596	3.10
simplebooks-92	0.00	10.67	413.45	0.626	—
simplebooks-92	0.02	10.75	127.30	0.554	4.77
simplebooks-92	0.05	10.67	130.91	0.557	4.88

Table 1: Regularization metrics at the three most relevant margin settings. Full two-dataset results for all five trained values are reported in Appendix E.2 .

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure E.8 : Robustness on simplebooks-92 : relative degradation under unstructured Gaussian noise and structured Gaussian drift. (a) Unstructured Gaussian noise: grouped boxplots of perturbed perplexity / clean perplexity across 100 validation subsamples for $\sigma\in\{0,0.05,0.10,0.15,0.20,0.25\}$ . (b) Structured Gaussian drift: grouped boxplots for $\sigma\in\{0,1.5,3.0,4.5,6.0\}$ . The same qualitative pattern as in the main text appears: a strong advantage for margin regularization under both perturbation families.

Level	Cross-entropy-only	Margin-regularized
Clean	1.00	1.00
$\sigma{=}0.05$	51.10	55.75
[95%]	[44.71, 62.41]	[48.01, 70.16]
$\sigma{=}0.10$	311.32	116.99
[95%]	[271.24, 401.72]	[99.58, 152.81]
$\sigma{=}0.15$	525.27	180.72

Appendix

Table 2: Unstructured Gaussian noise robustness on wikitext-103 . Entries report the median perplexity ratio (perturbed perplexity / clean perplexity); each estimate is followed by its empirical 95% interval on the line below.

Level	Cross-entropy-only	Margin-regularized
Clean	1.00	1.00
$\sigma{=}1.50$	59.48	50.88
[95%]	[38.07, 86.18]	[35.12, 83.29]
$\sigma{=}3.00$	201.56	135.75
[95%]	[131.63, 312.16]	[86.22, 219.54]
$\sigma{=}4.50$	329.51	207.99

Appendix

Table 3: Structured Gaussian drift robustness on wikitext-103 . Entries report the median perplexity ratio (perturbed perplexity / clean perplexity); each estimate is followed by its empirical 95% interval on the line below.

Level	Cross-entropy-only	Margin-regularized
Clean	1.00	1.00
$\sigma{=}0.05$	33.41	32.82
[95%]	[29.32, 36.88]	[29.71, 37.11]
$\sigma{=}0.10$	212.11	75.98
[95%]	[187.51, 233.99]	[67.99, 85.44]
$\sigma{=}0.15$	338.38	90.81

Appendix

Table 4: Unstructured Gaussian noise robustness on simplebooks-92 . Entries report the median perplexity ratio (perturbed perplexity / clean perplexity); each estimate is followed by its empirical 95% interval on the line below.

Level	Cross-entropy-only	Margin-regularized
Clean	1.00	1.00
$\sigma{=}1.50$	41.05	23.25
[95%]	[31.08, 58.49]	[18.62, 30.14]
$\sigma{=}3.00$	127.55	56.74
[95%]	[84.59, 214.33]	[44.73, 74.50]
$\sigma{=}4.50$	202.22	83.00

Appendix

Table 5: Structured Gaussian drift robustness on simplebooks-92 . Entries report the median perplexity ratio (perturbed perplexity / clean perplexity); each estimate is followed by its empirical 95% interval on the line below.

Figure E.9 : Selecting $\lambda_{m}$ on simplebooks-92 . Grouped boxplots over 100 validation subsamples for $\lambda_{m}\in\{0.00,0.02,0.05,0.10,0.25\}$ . (a) Degradation ratio at fixed unstructured Gaussian noise $\sigma{=}0.25$ . (b) Contextual embedding error. (c) Effective support size.

$\lambda_{m}$	Clean perplexity	Degradation ratio	Contextual embedding error	Effective support size
0.00	15.68	775.81	0.608	—
[95%]	[12.40, 18.10]	[661.68, 974.21]	[0.567, 0.630]	—
0.02	15.74	328.51	0.583	2.75
[95%]	[12.53, 18.15]	[273.81, 424.92]	[0.550, 0.614]	[2.06, 3.77]
0.05	15.71	364.92	0.596	3.10
[95%]	[12.50, 18.12]	[308.30, 473.11]	[0.563, 0.625]	[2.13, 4.02]

Appendix

Table 6: Regularization metrics on wikitext-103 . Degradation ratio is the perturbed perplexity / clean perplexity at fixed unstructured Gaussian noise

\sigma{=}0.25

. Contextual embedding error is the normalized squared distance between the true token embedding and the embedding predicted from the model’s next-token distribution given the preceding context. Effective support size is the entropy-based effective number of support tokens under the learned prior. It is not reported at

\lambda_{m}{=}0.00

because no prior is trained there. Each estimate is followed by its empirical 95% interval on the line below.

$\lambda_{m}$	Clean perplexity	Degradation ratio	Contextual embedding error	Effective support size
0.00	10.67	413.45	0.626	—
[95%]	[9.76, 11.76]	[372.86, 462.31]	[0.610, 0.641]	—
0.02	10.75	127.30	0.554	4.77
[95%]	[9.82, 11.71]	[114.70, 142.45]	[0.532, 0.580]	[3.83, 5.85]
0.05	10.67	130.91	0.557	4.88
[95%]	[9.81, 11.67]	[118.78, 148.52]	[0.534, 0.586]	[3.83, 6.15]

Appendix

Table 7: Regularization metrics on simplebooks-92 . Degradation ratio, contextual embedding error, and effective support size are defined as in Table 6 . Each estimate is followed by its empirical 95% interval on the line below.

Self-Attention

Self-attention lets sequence elements exchange context through learned query-key matches, powering Transformers and their efficient variants.

See all 225 papers

Stochasticity in Tokenisation Improves Robustness
Apr 17, 2026Sophie Steger, Rui Li, Sofiane Ennadir +4Adversarial Attacks on LLMs Neural Network Robustness
Institute of Signal Processing and Speech Communication, Graz University of Technology, Graz, Austria · ELLIS Institute Finland & Aalto University, Espoo, Finland · King AI Labs, Microsoft Gaming +2
Probabilistic Attribution For Large Language Models
May 20, 2026Shilpika Shilpika, Carlo Graziani, Bethany Lusch +2LLM Interpretability LLM Uncertainty Estimation
Argonne Leadership Computing Facility Argonne National Laboratory Lemont, IL, USA · Mathematics and Computer Science Division Argonne National Laboratory Lemont, IL, USA · Department of Computer Science University of Illinois Chicago Chicago, IL, USA
Long Context Pre-Training with Lighthouse Attention
May 7, 2026Bowen Peng, Subho Ghosh, Jeffrey QuesnelleAttention Mechanisms Self-Attention
Nous Research

Apr 17, 2026J/K move · Enter open · S save

Stochasticity in Tokenisation Improves Robustness

Sophie Steger, Rui Li, Sofiane Ennadir, Anya Sims+3

Institute of Signal Processing and Speech Communication, Graz University of Technology, Graz, Austria · ELLIS Institute Finland & Aalto University, Espoo, Finland · King AI Labs, Microsoft Gaming+2

Adversarial Attacks on LLMs Neural Network Robustness Adversarial Robustness of Language Models Subword Tokenization

Open paper AskPDF arXiv

The widespread adoption of large language models (LLMs) has increased concerns about their robustness. Vulnerabilities in perturbations of tokenisation of the input indicate that models trained with a deterministic canonical tokenisation can be brittle to adversarial attacks. Recent studies suggest that stochastic tokenisation can deliver internal representations that are less sensitive to perturbations. In this paper, we analyse how stochastic tokenisations affect robustness to adversarial attacks and random perturbations. We systematically study this over a range of learning regimes (pre-training, supervised fine-tuning, and in-context learning), data sets, and model architectures. We show that pre-training and fine-tuning with uniformly sampled stochastic tokenisations improve robustness to random and adversarial perturbations. Evaluating on uniformly sampled non-canonical tokenisations reduces the accuracy of a canonically trained Llama-1b model by 29.8%. We find that training with stochastic tokenisation preserves accuracy without increasing inference cost.

Abstract

Self-Attention

Explore similar work

Stochasticity in Tokenisation Improves Robustness

Probabilistic Attribution For Large Language Models

Long Context Pre-Training with Lighthouse Attention