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 x3 given x1=0 , x2=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−aVart (Eq. 8 ) evaluated across x3 . (c) Density profiles for varying coupling strengths a∈{−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 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 Vart . 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, λ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 σ∈{0,0.05,0.10,0.15,0.20,0.25} . (b) Structured Gaussian drift: grouped boxplots for σ∈{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 bt , 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 λm on wikitext-103 . Grouped boxplots over 100 validation subsamples for λm∈{0.00,0.02,0.05,0.10,0.25} . (a) Degradation ratio at fixed unstructured Gaussian noise σ=0.25 . (b) Contextual embedding error. (c) Effective support size.
Dataset
λ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 σ∈{0,0.05,0.10,0.15,0.20,0.25} . (b) Structured Gaussian drift: grouped boxplots for σ∈{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
σ=0.05
51.10
55.75
[95%]
[44.71, 62.41]
[48.01, 70.16]
σ=0.10
311.32
116.99
[95%]
[271.24, 401.72]
[99.58, 152.81]
σ=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
σ=1.50
59.48
50.88
[95%]
[38.07, 86.18]
[35.12, 83.29]
σ=3.00
201.56
135.75
[95%]
[131.63, 312.16]
[86.22, 219.54]
σ=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
σ=0.05
33.41
32.82
[95%]
[29.32, 36.88]
[29.71, 37.11]
σ=0.10
212.11
75.98
[95%]
[187.51, 233.99]
[67.99, 85.44]
σ=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
σ=1.50
41.05
23.25
[95%]
[31.08, 58.49]
[18.62, 30.14]
σ=3.00
127.55
56.74
[95%]
[84.59, 214.33]
[44.73, 74.50]
σ=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 λm on simplebooks-92 . Grouped boxplots over 100 validation subsamples for λm∈{0.00,0.02,0.05,0.10,0.25} . (a) Degradation ratio at fixed unstructured Gaussian noise σ=0.25 . (b) Contextual embedding error. (c) Effective support size.
λ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 σ=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 λm=0.00 because no prior is trained there. Each estimate is followed by its empirical 95% interval on the line below.
λ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.
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.
Sophie Steger, Rui Li, Sofiane Ennadir +4
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
The generative nature of Large Language Models (LLMs) is reflected in the conditional probabilities they compute to sample each response token given the previous tokens. These probabilities encode the distributional structure that the model learns in training and exploits in inference. In this work, we use these probabilities to situate LLMs within the mathematical theory of stochastic processes. We use this framework to design a model-agnostic probabilistic token attribution measure, using Bayes rule to invert the next-token log-probabilities so as to capture the models internal representation of the distribution over token sequences. The representation is independent of the models computational structure. This representation yields the conditional probability of the response given the prompt, and of the response given the prompt with a token marginalized away. Our attribution score is the log of the ratio of these probabilities. We further compute the entropies of a single prompts token distributions, conditioned on the remaining context. The interplay between entropy and attribution score sheds light on LLM behavior. We evaluate 8 models across 7 prompts and investigate anomalies, token sensitivity, response stability, model stability, and training convergence, thereby improving interpretability and guiding users to focus on uncertain or unstable parts of the generation.
Shilpika Shilpika, Carlo Graziani, Bethany Lusch +2
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
Training causal transformers at extreme sequence lengths is bottlenecked by the quadratic time and memory of scaled dot-product attention (SDPA). In this work, we propose Lighthouse Attention, a training-only symmetrical selection-based hierarchical attention algorithm that wraps around ordinary SDPA and can be easily removed towards the end of the training. Our hierarchical selection is also gradient-free, which exempts us from dealing with a complicated and potentially inefficient backward pass kernel. Our contribution is three-fold: (i) A subquadratic hierarchical pre- and post-processing step that does adaptive compression and decompression of the sequence. (ii) A symmetrical compression strategy that pools queries, keys and values at the same time, while preserving left-to-right causality, which greatly improves parallelism. (iii) A two stage training approach which we pre-train for the majority of the time with Lighthouse Attention and recover a full attention model at the end with a short training. We run preliminary small scale LLM pre-training experiments that show the effectiveness of our method compared to full attention training with all other settings matched, where we achieve a faster total training time and lower final loss after the recovery phase. Full code is available at: https://github.com/ighoshsubho/lighthouse-attention