PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashAttention with Statistics: Low-Overhead Attention Features for Compact Hallucination Detectors

Andrei Volodichev, Maksim Savichev, Nikita Sviridov, Artem Yuliy, Dmitry Simakov, Alexey Zaytsev

arXiv:2610.21483v1cs.AI

TL;DR

Attention-based hallucination detectors use proxy-model attention statistics, but standard pipelines materialize an attention matrix whose memory grows quadratically with sequence length. This paper computes the statistics through FlashAttention-compatible paths and evaluates extraction cost and frozen-proxy detectors. The proxies perform strongly across three benchmarks, while attention features significantly improve F1 in half of the benchmark–proxy pairs without significant declines.

  • Problem

    Attention-based detectors use summary statistics of proxy attention weights, but standard pipelines materialize the full attention matrix even though detectors keep only a linear number of statistics.

    Source: §1 · p. 1

  • Method

    The paper computes attention statistics during or after a FlashAttention pass using fused accumulation, post-pass lookups, or a second pass over scores, and evaluates frozen Qwen3 proxy detectors.

    Source: §1 · pp. 1–2

  • Results

    A 4B proxy outperforms the strongest evaluated LLM judge in F1 on two benchmarks, and adding attention features significantly improves F1 in 6 of 12 benchmark–proxy pairs without significantly lowering it.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Across all 12 dataset–proxy pairs, A64H64 achieves higher F1 and AUROC than prompting the corresponding Qwen3 model as a judge.

    Source: §5.4 · p. 8

    The implementation is not yet a production kernel and does not support sliding-window or soft-capped attention or arbitrary head dimensions.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Large language models can produce fluent answers that the supplied context does not support. A compact frozen proxy model can detect such answers by reading the context and another model's response, even from a proprietary generator. Attention-based detectors use summary statistics of the proxy's attention weights, but standard pipelines materialize the attention matrix, whose memory grows quadratically with sequence length. We compute these statistics without materializing attention, by accumulation in the FlashAttention pass, sparse lookups after it, or a second tiled pass over the scores, and classify them by this path into nearly free and costlier ones. On an H100, the statistics of five published detectors add 16% to the latency of a Qwen3-8B forward at 2K tokens, falling to 1.4% at 35K, and store 22.5 KiB per response token. On three benchmarks, probes trained on frozen Qwen3 proxies of 0.6B to 8B parameters prove strong detectors: a 4B proxy outperforms the strongest of nine zero-shot 4-9B LLM judges in F1 on two of the three, the proxies match or outperform specialized grounding checkers of up to 7B parameters, probing each proxy beats prompting it as a judge, and adding attention features significantly improves F1 in half of the 12 benchmark-proxy pairs, never significantly lowering it.

1 Introduction

Contextual hallucination detectors can use a compact proxy to assess whether another model’s answer is supported, but attention-based approaches typically materialize a large attention matrix. This paper computes attention statistics without materializing that matrix and evaluates both their extraction cost and proxy-detector performance.

1 Introduction

  • Standard attention-statistics pipelines materialize the attention matrix, which can occupy about 1 GB at 4K tokens and tens of GB at 16K–32K for a 32-head layer.Detectors retain only a linear number of statistics.Source: §1 · p. 1
  • 16% at 2K tokens and 1.4% at 35K tokens: extracting statistics for five published detectors adds this much latency to a Qwen3-8B forward.The implementation stores 22.5 KiB per response token.Source: §1 · p. 2
  • The implementation computes statistics by accumulation during FlashAttention, post-pass lookups, a second pass over scores, or dense materialization, and classifies them by computation path.The authors implement statistics for five published detectors in Triton.Source: §1 · pp. 1–2
  • A 4B proxy achieves higher observed F1 than the strongest evaluated 4–9B LLM judge on two benchmarks, while a 0.6B proxy does so on the third.The 0.6B proxy trails that judge by at most 0.03 F1 on the other two benchmarks.Source: §1 · p. 2
  • Adding attention features significantly improves F1 in 6 of 12 benchmark–proxy pairs and never significantly lowers it.Probes use frozen Qwen3 proxies ranging from 0.6B to 8B parameters.Source: §1 · p. 2

2 Background

The paper defines attention statistics and features for a teacher-forced detector that reads response rows, then describes how FlashAttention computes exact attention in tiles without storing the weights.

2.1 Setting and notation

  • A statistic is a per-head row or column reduction of attention weights; a feature is a classifier input derived from statistics, and a detector recipe selects features, heads, and layers.The response is known in advance: it follows a prompt of P tokens, and the targeted detectors read statistics from response rows.Source: §2.1 · p. 3
  • Under the causal mask, prompt rows depend only on prompt keys, so any production kernel can compute them; the response stage evaluates response queries against all keys and values.The response block has R × L scores and weights, with NR = RP + R(R + 1)/2 unmasked scores per head without padding.Source: §2.1 · p. 3

2.2 Online softmax in FlashAttention

  • FlashAttention-2 evaluates the response block in a query-major tiled pass, keeping each query tile resident while key tiles stream past, without storing attention weights.The paper calls one tiled evaluation of every response-block score a pass.Source: §2.2 · p. 3
  • Online softmax carries a running maximum, normalizer, and unnormalized output; after the last tile, FlashAttention computes the output and stores the log-sum-exp.A later kernel can recover a weight by re-evaluating its score and using the stored log-sum-exp.Source: §2.2 · p. 3

3 Computing Attention Statistics without Materialization

The paper classifies statistics by the extra score evaluations and storage their computation requires, then describes fused accumulation and post-pass paths for computing them without materializing attention.

3.1 Computation paths

  • T1 through T4 require, respectively, no extra scores, a fixed number of lookups per row, a second evaluation of every score, or materialization of the response block.The tiers describe the paths used by the implementation and bound their cost from above; the stored log-sum-exp belongs to no tier.Source: §3.1 · p. 3

3.2 T1 (fused): statistics in online softmax

  • T1 accumulates statistics inside the attention pass using bounded state, without extra scores or exponentials, and with state size independent of sequence length.The accumulator can compute fixed functions of the resulting reductions and row log-sum-exp.Source: §3.2 · p. 4
  • Prompt mass, entropy, effective support, and maximum weight are examples of T1 statistics computed by updating accumulators alongside FlashAttention’s state.For prompt mass, the kernel accumulates exponentials for prompt keys and divides by the final normalizer.Source: §3.2 · p. 4

3.3 T2–T4: post-pass statistics

  • T2 computes selected weights such as diagonal, previous-token, or fixed sink-column weights by re-evaluating a fixed number of scores per response row and using the stored log-sum-exp.These extra scores grow linearly with sequence length; the chosen post-pass kernel can follow any forward kernel that stores the log-sum-exp.Source: §3.3 · p. 4
  • T3 re-evaluates all response-stage scores for column reductions or threshold counts, while T4 materializes the response block for statistics that need the dense map.T3 uses a key-major loop and costs close to a second response stage; T4 has O(NR) memory and no implemented kernel.Source: §3.3 · p. 4

4 Kernel Benchmark

The benchmark measures the accuracy and cost of extracting attention statistics with tiled kernels, including comparisons with materialized attention and production kernels. T1 and T2 add modest overhead, while T3 costs more; on real inputs the tiled approach uses substantially less memory than dense extraction.

Benchmark setup

  • For Qwen3-8B shapes at d = 128 and r = 0.3, latency is measured from 1K to 100K tokens, with overhead also compared against whole-sequence SDPA cuDNN.The benchmarks use one H100 NVL; configurations run response rows through Triton kernels and prompt rows through PyTorch SDPA.Source: Figure 3 · p. 5, §4 · p. 5

Numerical accuracy

  • Our kernels produce statistics three to four orders of magnitude more accurately than materialized attention for every statistic tested.Against an fp64 reference, the kernels reduce fp32 probabilities in registers, while materialized attention rounds each probability to bf16 first.Source: §4 · p. 5

Marginal cost per tier

  • At most, ≤T2 adds 5 percentage points of compute time over ≤T1; both add 1–30% over Base, while ≤T3 adds 16–87%.The tier costs were measured across response fractions and head dimensions; ≤T3's second score pass adds matrix work.Source: §4 · p. 5

Scaling and tier cost

  • At 100K tokens and r = 0.3, a layer with all row statistics takes 387 ms, versus 351 ms for cuDNN and 523 ms for FA2 without statistics.The kernels scale to 100K tokens without an expected wall; below about 4K tokens, fixed CPU overhead dominates.Source: §4 · pp. 5–6, Figure 3 · p. 5

Real-input comparison

  • On the sample with 785 response tokens in 7.2K, dense extraction adds 501 ms and 2.1 GiB, while the kernels add 53 ms and 23 MiB.Across the evaluated prompts, dense latency starts growing once responses exceed about 100 tokens, while dense memory grows with RL and kernel memory with R.Source: §4 · p. 6, Figure 4 · p. 6

5 Proxy Detection

The paper evaluates frozen Qwen3 proxies as detectors of whether an answer is supported by its context, using hidden states, logits, and attention-derived features. Across three benchmarks, proxy detectors compare favorably with prompted judges, while the value of attention features and proxy size varies by dataset.

5.1 Task, data, and metrics

  • Each example is read in one teacher-forced forward by a frozen 0.6B, 1.7B, 4B, or 8B Qwen3 proxy; the generator need not be the proxy.The detector predicts whether the context supports an existing answer; evaluation uses RAGTruth QA, HaluBench, and TRIVIA+ with related examples kept within splits.Source: §5.1 · p. 6

5.2 Feature blocks

  • Feature blocks comprise hidden states and logits from the proxy forward and attention variants A7 and A64; per-token quantities are pooled over answer tokens using mean, standard deviation, minimum, and maximum.A7 uses selected induction-head statistics, while A64 uses each head's largest sink values from T3, projected to 64 principal components.Source: §5.2 · p. 7

5.3 Selection and training per dataset and proxy

  • Feature sets are evaluated with class-balanced L2-regularized logistic regression, with scaling, projections, and hyperparameters fit or selected without test labels.A7 head count is selected by validation AUROC; A64 sink count is chosen by grouped five-fold cross-validation on training data.Source: §5.3 · p. 7

5.4 Main results

  • The best-F1 feature set contains A7 or A64 in all 12 dataset–proxy pairs, and the best-AUROC set does too; no tested addition significantly lowers F1.After Holm correction, A7 significantly improves F1 for all four proxies on RAGTruth QA with logits, and for two proxies with hidden states; A64 improves hidden-state F1 for the 4B and 8B proxies on HaluBench.Source: §5.4 · pp. 7–8
  • A64 alone beats A7 in all 12 pairs on F1 and AUROC, but with hidden states A64H64 wins on HaluBench and TRIVIA+ while A7H64 wins on RAGTruth QA for three of four proxies.Larger proxies improve hidden-state, logit, or all-head-sink detectors on HaluBench and TRIVIA+, while even the smallest proxy remains competitive on RAGTruth QA.Source: §5.4 · p. 8
  • A 4B proxy achieves higher observed F1 than Qwen3.5-9B, the strongest of nine evaluated judges, on HaluBench and TRIVIA+; a 0.6B proxy does so on RAGTruth QA.A64H64 achieves higher F1 and AUROC than prompting the corresponding Qwen3 model in all 12 dataset–proxy pairs.Source: §5.4 · p. 8

5.5 Cost of feature extraction

  • For Qwen3-8B, extracting all kernel statistics adds 16% latency at 2K tokens and 1.4% at 35K; the features add about 28 KiB per response token in peak memory.The measured statistics occupy 22.5 KiB per response token; short-input speed is limited by Python-side CPU overhead.Source: §5.5 · p. 8

6 Conclusion

The paper presents a FlashAttention-compatible backend that computes statistics underlying five detectors without materializing attention matrices. It reports reduced latency overhead at longer sequence lengths and finds that supervised detectors built on frozen proxies can compete with larger prompted judges.

Conclusion

  • On an H100 NVL at batch size one, statistics for five detectors add 16% to Qwen3-8B forward latency at 2K tokens and 1.4% at 35K.Reducing Python overhead and GPU synchronization would further improve extraction speed on short inputs.Source: §6 · p. 8
  • The backend computes attention statistics with memory linear in sequence length, without materializing attention matrices.Its implementation classifies statistics by whether they reuse the attention pass or require lookups or a second pass over scores.Source: §6 · p. 8
  • Supervised detectors built on frozen proxy models can compete with larger prompted judges.The conclusion frames the implementation as supporting attention features used by existing detectors and their use with small, frozen proxy models.Source: §6 · p. 8

Limitations

The reported tiers describe the paths implemented, not every efficient exact algorithm or an optimal implementation. The backend and proxy also have concrete performance, deployment, and evaluation boundaries.

Limitations

  • The tiers are implementation-specific and do not claim to enumerate all efficient statistics or optimal algorithms.A statistic assigned to T4 may still admit an efficient exact kernel the authors did not implement.Source: Limitations · p. 9
  • The Triton kernels may be slower than specialized CUDA kernels, and the Python wrapper adds 47–60 ms of fixed host cost per forward.The response-stage kernel is faster than FA2 on Hopper but slower than cuDNN; the host cost is mostly not GPU work.Source: Limitations · p. 9
  • The implementation is not production-ready: it processes a complete response at once and lacks support for several attention configurations and arbitrary head dimensions.It does not yet implement token-by-token generation, sliding-window or soft-capped attention, or dimensions beyond 64, 128, and 256.Source: Limitations · p. 9
  • The proxy needs an auxiliary forward pass and labeled probe-training examples, and a small proxy may miss evidence available to a larger generator.Evaluation concerns support relative to supplied context; dataset artifacts, label disagreements, truncation, and domain shift can affect results.Source: Limitations · p. 9

Ethical Considerations

The authors aim to make hallucination detection inexpensive enough to check every response in a retrieval-augmented system, while positioning it as a complement to human review in high-stakes settings.

Ethical Considerations

  • The authors intend the detector to complement, not replace, human review in high-stakes settings because it can miss unsupported answers or flag supported ones.The experiments use public benchmarks under their terms, with no new data collection or human subjects.Source: Ethical Considerations · p. 9

Guide to the appendices

The appendices provide methodological details, kernel benchmarks, proxy-detection results, and supplementary experiments. The related-work discussion connects efficient attention kernels with attention-based hallucination detectors by mapping detector statistics to matrix-free computation paths.

Guide to the appendices

  • The appendices cover statistic definitions and implementation, kernel benchmarks and memory, proxy-detection evaluations, cross-domain transfer, and judge prompts.They also include numerical accuracy, cost by tier and sequence length, and whole-model latency measurements.Source: Guide to the appendices · p. 12

A Related Work

  • The related work links fused attention reductions and attention-based detectors by mapping detector statistics to exact matrix-free paths, measuring their cost, and using them in proxy detectors.Prior work includes fused reductions and key-eviction scoring, while the cited hallucination detectors read attention statistics; related proxy moderation systems use their own attention rather than generating a verdict.Source: A Related Work · p. 12

B Statistics computed by the kernels

The kernels compute detector statistics on response rows through in-pass accumulators, fixed lookups, or a second score pass, then derive detector features from the returned values. Response-only computation changes two published features, and their adapted versions have not been compared in detection.

B Statistics computed by the kernels

  • Response-stage kernels return per-row or per-column statistics, which the host reduces over response rows and selected heads to form detector features.The host-side reductions take O(R) or O(L) time per head.Source: B Statistics computed by the kernels · pp. 12–13
  • Response-only averages change LLM-Check’s log-diagonal feature and SinkProbe’s per-key score; the paper has not compared these adapted versions in detection.Lookback Lens and RAUQ features are already defined on generated tokens.Source: B Statistics computed by the kernels · p. 13
  • BOS sink rates use T2 only when batch sequences share padding; other sink definitions, or unequal padding, fall back to T3.The normalized sink rate averages a key’s attention weight over response rows that can see it.Source: B Statistics computed by the kernels · p. 13

C Separable row reductions

  • For T1, each score tile updates statistic accumulators alongside FlashAttention, retaining no per-key state; the resulting values and LSE support derived features after the pass.Examples include prompt mass, entropy, effective support, and maximum weight.Source: D Implementation · p. 14, C Separable row reductions · p. 14

D Implementation

  • T2 looks up fixed positions using Q, K, and LSE, while T3 computes column sums with a key-major pass; sliding-window and soft-capped layers instead use default attention without statistics.T2 works after any forward kernel that writes LSE; T3 writes column sums without atomic additions.Source: D Implementation · p. 15

E Kernel benchmark setup

The kernel statistics are more accurate than materialized attention against an fp64 reference, while their layer-level overhead depends on the computation tier and response fraction. The benchmarks use an H100 NVL and compare the kernels with FlashAttention-2 and cuDNN.

E Kernel benchmark setup

  • The benchmark compares statistics-off and statistics-enabled attention on one H100 NVL in bf16, using PyTorch SDPA with FlashAttention-2 and cuDNN as production baselines.Timing medians use 100 repetitions after 25 warm-ups for configuration sweeps; the length-scaling runs use 15–80 repetitions.Source: E Kernel benchmark setup · p. 15

F Numerical accuracy (Q1)

  • Three to four orders of magnitude closer to fp64, the kernels have lower error for every statistic, input, and tested length than materialized attention.The comparison covers random and PHANTOM inputs at 2K, 8K, and 32K tokens.Source: Table 4 · p. 16
  • Our kernels’ per-head means rank heads like fp64, with Spearman correlation 1.000; the dense path reaches at least 0.998, and kernel entropy mean error stays below 4.7e−5 at 32K.For heads with H < 0.25 at 32K, dense entropy error reaches 1.4e−2.Source: F Numerical accuracy (Q1) · p. 16

G Cost per tier and length scaling (Q2)

  • At r = 0.3, ≤T2 adds 3–14% over Base, ≤T3 adds 32–47%, and Base takes 4–14% longer than cuDNN from 15K to 100K tokens.These overhead bands do not grow with sequence length over that range.Source: Figure 3b

H Whole-model latency and extraction cost (Q3)

Whole-model latency measurements score all layers of Qwen3-8B and Qwen3-0.6B on PHANTOM and TRIVIA+ inputs. On short inputs, feature extraction adds a fixed host-time cost rather than a cost that scales with input length.

H Whole-model latency and extraction cost (Q3)

  • The whole-model timing setup scores all heads and layers of Qwen3-8B and Qwen3-0.6B on PHANTOM and TRIVIA+ inputs, with bins containing fewer than 12 samples omitted.Prompt attention uses FlashAttention-2 up to 8K tokens and cuDNN for longer inputs; Q3 uses FlashAttention-2 throughout.Source: H Whole-model latency and extraction… · p. 17
  • Below about 600 tokens for Qwen3-8B and 3K for Qwen3-0.6B, features add fixed costs of 60 and 47 ms, respectively.That is about 1.7 ms per scored layer for both models, despite the smaller model’s layers doing a fraction of the work.Source: H Whole-model latency and extraction… · p. 17
  • Feature extraction overhead is host time: each scored layer launches 88 more kernels and reads ten values back from the GPU, stalling the host until the queue drains.Profiling attributes 0.8 ms per layer to launching the three statistics kernels and another 0.8 ms to assembling features; copying layer configurations adds 9 ms per forward.Source: H Whole-model latency and extraction… · p. 17
  • A split layer called on its own pays about 0.3–0.6 ms of host time for setup and 19 kernel launches across all row statistics.This is distinct from the whole-forward short-input overhead, which includes per-layer wrapper and feature-assembly costs.Source: H Whole-model latency and extraction… · p. 18

I Memory

Whole-forward feature memory grows with response length rather than prompt length, while measured layer-level peaks and short-input latency and memory depend on execution setup. The per-layer comparison finds close feature agreement with dense materialization.

I Memory

  • At 32K, a layer's 520–532 MiB extra peak includes 7.6–20 MiB of statistics; concatenating its two stage outputs explains most of the difference from SDPA.The same layer's attention matrix would occupy 64 GiB.Source: I Memory · p. 18
  • 26–32 KiB per response token is the measured whole-forward peak growth, mostly from the 22.5 KiB of Qwen3-8B statistics stored per response token.The fit found growth with response length, not prompt length; overhead was at most 3.2% on PHANTOM and 6% median, up to 17% on short RAGTruth QA inputs.Source: I Memory · p. 18
  • At 35K tokens, Figure 6's measured peak grows about 28 KiB per response token for Qwen3-8B and 11 KiB for Qwen3-0.6B.The features themselves store 22.5 and 8.75 KiB per response token, respectively.Source: I Memory · p. 18, Figure 6 · p. 19

J Per-layer comparison on short inputs

  • At batch 1, Ours is slower than Dense (1.26 vs 1.19 ms) but uses less peak memory (14 vs 116 MiB); at batch 8, it is faster (4.71 vs 7.60 ms) and uses 136 MiB vs 1.4 GiB.These are per-layer calls on short real inputs; memory is the PyTorch allocated peak above memory allocated before the call, excluding weights for whole-forward measurements.Source: I Memory · p. 18, J Per-layer comparison on short inputs · p. 19
  • Ours differs from Dense by mean relative error from 0.04% on LLM-Check's log-diagonal feature to 1.6% on previous-token weight; effective-support mean error against fp64 is 5.9e−3 for Ours and 5.3 for Dense.The largest absolute difference, 6.8, is on effective support, whose typical values are about 12.Source: J Per-layer comparison on short inputs · p. 19

K Data

This section documents benchmark construction, feature and evaluation protocols, detector and transfer findings, and judge-prompting criteria, alongside artifact licensing and AI-assisted proofreading.

K Data

  • On unfiltered HaluBench, answer-only TF–IDF reaches 0.842 AUROC, while length alone reaches 0.586 overall after filtering balances answer length.The filtering retains 1,108 HaluEval answers and 6,008 answers overall; groups are assigned to splits to approach 70/15/15% within each source and label, and remain intact across splits.Source: K Data · pp. 19–20, Table 8 · p. 20

L.1 Feature recipes

  • Feature recipes were selected using Qwen3-4B cross-validation on TRIVIA+ training data and reused; they combine induction-ranked A7 attention features, A64 sink values, H64 hidden states, and L9 logit summaries.A7 keeps seven attention features; A64 projects per-head sink values onto 64 components, while H64 pools hidden-state means and standard deviations and projects them onto 64 components.Source: §L.1 · pp. 20–21

L.2 Logit features

  • The logit block computes seven per-token statistics from the proxy’s next-token distribution and pools them over answer tokens into 224 candidate features.The output head is applied only at positions predicting answer tokens and the preceding position, avoiding a full vocabulary distribution for the prompt.Source: §L.2 · p. 21

L.3 Training and significance tests

  • Feature-set models select regularization and thresholds on validation data; attention-block additions are tested with paired group bootstraps and Holm correction across 12 dataset–proxy pairs.The tests use 100,000 resamples with models, regularization, and thresholds fixed; only large gains remain significant given 150–546 test groups.Source: §L.3 · pp. 21–22

L.4 Full results

  • No tested attention-feature addition significantly lowers F1, and significant gains occur where each block helps most.Tables 11 and 12 test adding A7 and adding A64 to H64, respectively; Table 10 reports test AUROC for the detectors.Source: §L.4 · p. 22, Table 11 · p. 23, Table 12 · p. 23

L.5 LLM judges

  • LLM judges receive the context, question, and answer, return a True/False hallucination verdict, and use validation to select prompts or thresholds.The judge templates use either no demonstrations or eight labeled examples; Appendix N includes strict prompts requiring a one-word verdict and judging only against the supplied context.Source: §L.5 · p. 22, N LLM judge prompts · p. 27

M Cross-domain transfer on HaluBench

  • RAUQ and INTRYGUE transfer best across HaluBench sources, losing at most 0.01 AUROC from fitting to other sources; Lookback Lens loses 0.15 and falls to chance (0.51).INTRYGUE remains at 0.38–0.41 on DROP, and no method exceeds 0.53 on FinanceBench; fitting on more sources helps little.Source: M Cross-domain transfer on HaluBench · p. 25, Table 14 · p. 26

N LLM judge prompts

  • The judge marks unsupported or incorrect claims as hallucinations, while treating supported answers, omissions, faithful paraphrases, and appropriate refusals as faithful.A single unsupported or incorrect claim suffices; plausible facts not supported by the context still count as hallucinations.Source: N LLM judge prompts · pp. 27–28

O Use of scientific artifacts

Related papers

Loading 2610.21483v1…