AP-OOD: Attention Pooling for Out-of-Distribution Detection

Out-of-distribution (OOD) detection, which maps high-dimensional data into a scalar OOD score, is critical for the reliable deployment of machine learning models. A key challenge in recent research is how to effectively leverage and aggregate token embeddings from language models to obtain the OOD score. In this work, we propose AP-OOD, a novel OOD detection method for natural language that goes beyond simple average-based aggregation by exploiting token-level information. AP-OOD is a semi-supervised approach that flexibly interpolates between unsupervised and supervised settings, enabling the use of limited auxiliary outlier data. Empirically, AP-OOD sets a new state of the art in OOD detection for text: in the unsupervised setting, it reduces the FPR95 (false positive rate at 95% true positives) from 27.84% to 4.67% on XSUM summarization, and from 77.08% to 70.37% on WMT15 En-Fr translation.

💡 Research Summary

The paper introduces AP‑OOD, a novel out‑of‑distribution (OOD) detection method tailored for natural‑language processing that moves beyond the simplistic mean‑pooling of token embeddings. The authors first point out that averaging token representations discards crucial token‑level structure, causing ID (in‑distribution) and OOD sequences to become indistinguishable when their means coincide. To address this, they reformulate the Mahalanobis distance as a directional decomposition and replace the fixed linear projection with a learnable attention‑pooling operation. Specifically, a query vector w (and optionally multiple queries W) and an inverse‑temperature β are introduced, yielding an attention‑pooled representation AttPoolβ(Z,w)=Z·softmax(β Zᵀw). This operation assigns adaptive weights to each token, preserving discriminative information that mean‑pooling loses.

The method defines a global prototype μ by applying the same attention pooling to a concatenated corpus‑wide matrix \tilde Z. The squared distance between a sequence representation Z and the global prototype is then computed as the sum over heads of the squared differences between attention‑weighted projections of Z and \tilde Z. When β=0 and the number of heads M equals the embedding dimension D, this distance reduces exactly to the classic Mahalanobis distance, establishing AP‑OOD as a strict generalization.

AP‑OOD operates in two regimes. In the unsupervised setting, only ID data are used; the loss minimizes the attention‑based distance, encouraging the model to learn typical token‑level patterns. In the supervised (or semi‑supervised) setting, a limited auxiliary outlier set (AUX) is incorporated. The loss adds a binary cross‑entropy term that pushes the distance for AUX samples upward while still minimizing it for ID samples. A weighting parameter λ smoothly interpolates between the two extremes, allowing practitioners to exploit any amount of outlier data without a hard switch.

Implementation details include a mini‑batch attention scheme to keep memory usage tractable, and extensions to multiple heads and multiple queries per head, which enable the model to capture diverse token‑level anomalies. A regularization term penalizes the norm of the query vectors to avoid degenerate solutions.

Experiments are conducted on two generation tasks: summarization (PEGASUS‑LARGE fine‑tuned on XSUM) and machine translation (WMT15 English‑French). OOD benchmarks comprise several out‑of‑domain corpora (CNN/DailyMail, Newsroom, Reddit, Samsum) and a distinct language pair for translation. In the unsupervised regime, AP‑OOD dramatically reduces the false‑positive rate at 95 % true‑positive recall (FPR95) from 27.84 % to 4.67 % on XSUM and from 77.08 % to 70.37 % on WMT15, while achieving the highest area under the ROC curve (AUR‑OC) across all baselines (Mahalanobis, K‑NN, Deep SVDD, perplexity, entropy). In the supervised regime, even with a modest AUX set, AP‑OOD consistently outperforms the same baselines, demonstrating graceful scaling with the amount of outlier supervision.

Ablation studies explore the impact of the temperature β, the number of heads M, and the regularization strength, confirming that the attention mechanism is the primary driver of performance gains. Visualizations illustrate how attention weights focus on tokens that are atypical for the ID distribution, and theoretical appendices prove the equivalence to Mahalanobis distance and provide a kernel‑view interpretation of the attention pooling.

Limitations are acknowledged: constructing the global concatenated matrix \tilde Z can be memory‑intensive for very long sequences, and hyper‑parameter sensitivity may require automated tuning. Future work could extend AP‑OOD to decoder‑only models (e.g., GPT‑style) and investigate streaming or approximate attention mechanisms for large‑scale corpora.

In summary, AP‑OOD offers a principled, flexible, and empirically strong solution for OOD detection in text generation models. By preserving token‑level information through learnable attention pooling and supporting both unsupervised and semi‑supervised scenarios, it sets a new state of the art and provides a practical tool for deploying safer language models.