# Uncertainty Awareness of Large Language Models Under Code Distribution Shifts: A Benchmark Study

arXiv: https://arxiv.org/abs/2402.05939 · HTML: https://mlsanity.com/paper/2402.05939/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 43 references that are in the map.
- First posted on arXiv 12 Jan 2024.

## Builds on (most cited first, 29 of 29)

- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Large Language Models are Zero-Shot Reasoners](https://mlsanity.com/paper/2205.11916/) (2022)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [A Baseline for Detecting Misclassified and Out-of-Distribution Examples in Neural Networks](https://mlsanity.com/paper/1610.02136/) (2016)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [A Simple Unified Framework for Detecting Out-of-Distribution Samples and Adversarial Attacks](https://mlsanity.com/paper/1807.03888/) (2018)
- [Enhancing The Reliability of Out-of-distribution Image Detection in Neural Networks](https://mlsanity.com/paper/1706.02690/) (2017)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [Deep Bayesian Active Learning with Image Data](https://mlsanity.com/paper/1703.02910/) (2017)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness](https://mlsanity.com/paper/2006.10108/) (2020)
- [code2vec: Learning Distributed Representations of Code](https://mlsanity.com/paper/1803.09473/) (2018)
- [Out-of-Distribution Detection and Selective Generation for Conditional Language Models](https://mlsanity.com/paper/2209.15558/) (2022)
- [code2seq: Generating Sequences from Structured Representations of Code](https://mlsanity.com/paper/1808.01400/) (2018)
- [Understanding Measures of Uncertainty for Adversarial Example Detection](https://mlsanity.com/paper/1803.08533/) (2018)
- [Uncertainty Quantification in Scientific Machine Learning: Methods, Metrics, and Comparisons](https://mlsanity.com/paper/2201.07766/) (2022)
- [Plex: Towards Reliability using Pretrained Large Model Extensions](https://mlsanity.com/paper/2207.07411/) (2022)
- [Contrastive Code Representation Learning](https://mlsanity.com/paper/2007.04973/) (2020)
- [Code Prediction by Feeding Trees to Transformers](https://mlsanity.com/paper/2003.13848/) (2020)
- [Adversarial Sample Detection for Deep Neural Network through Model Mutation Testing](https://mlsanity.com/paper/1812.05793/) (2018)
- [On the Generalizability of Neural Program Models with respect to Semantic-Preserving Program Transformations](https://mlsanity.com/paper/2008.01566/) (2020)
- [Self-Ensemble Adversarial Training for Improved Robustness](https://mlsanity.com/paper/2203.09678/) (2022)
- [Sequence Level Contrastive Learning for Text Summarization](https://mlsanity.com/paper/2109.03481/) (2021)
- [On Uncertainty Calibration and Selective Generation in Probabilistic Neural Summarization: A Benchmark Study](https://mlsanity.com/paper/2304.08653/) (2023)
- [Impact of Evaluation Methodologies on Code Summarization](https://mlsanity.com/paper/2108.09619/) (2021)
- [Uncertainty-Aware Bootstrap Learning for Joint Extraction on Distantly-Supervised Data](https://mlsanity.com/paper/2305.03827/) (2023)
- [CodeS: Towards Code Model Generalization Under Distribution Shift](https://mlsanity.com/paper/2206.05480/) (2022)

## Built on by (most cited first, 3 of 3)

- [UBench: Benchmarking Uncertainty in Large Language Models with Multiple Choice Questions](https://mlsanity.com/paper/2406.12784/) (2024)
- [LeMix: Unified Scheduling for LLM Training and Inference on Multi-GPU Systems](https://mlsanity.com/paper/2507.21276/) (2025)
- [PromptShift-CRC: Drift-Aware Conformal Risk Control for Foundation Models Under Prompt and Domain Shift](https://mlsanity.com/paper/2606.15964/) (2026)

## Abstract (arXiv)

Large Language Models (LLMs) have been widely employed in programming language analysis to enhance human productivity. Yet, their reliability can be compromised by various code distribution shifts, leading to inconsistent outputs. While probabilistic methods are known to mitigate such impact through uncertainty calibration and estimation, their efficacy in the language domain remains underexplored compared to their application in image-based tasks. In this work, we first introduce a large-scale benchmark dataset, incorporating three realistic patterns of code distribution shifts at varying intensities. Then we thoroughly investigate state-of-the-art probabilistic methods applied to CodeLlama using these shifted code snippets. We observe that these methods generally improve the uncertainty awareness of CodeLlama, with increased calibration quality and higher uncertainty estimation~(UE) precision. However, our study further reveals varied performance dynamics across different criteria (e.g., calibration error vs misclassification detection) and trade-off between efficacy and efficiency, highlighting necessary methodological selection tailored to specific contexts.