Adaptive Inference Under Distribution Shift: Conditional Diffusion, Test-Time Compute, and Selective Attention

Abstract

Distribution shift manifests differently across machine learning domains, shaping both the challenge and the solution in each case. In wireless channel estimation, the shift is continuous and physical: user mobility induces time-varying Doppler dynamics that evolve channel statistics on the timescale of seconds. In mathematical reasoning, shift is discrete and semantic: problem difficulty varies across instances in ways not observable from the input, so a system cannot know in advance how much computation a problem warrants. In medical imaging, shift is institutional and acquisitional: heterogeneity in scanner hardware and patient populations produces covariate shift, with pathological patterns appearing at varying scales and locations. Despite these differences, a common failure mode unites the three settings: static inference that commits to fixed computational strategies wastes resources when conditions are favorable and underperforms when they are demanding. This thesis develops adaptive inference mechanisms built on a shared principle of conditioning computational allocation on input-dependent signals, instantiated differently in each domain: for wireless channels as signal-to-noise ratio (SNR) combined with a learned temporal context; for large language model (LLM) reasoning as an intermediate quality signal from a learned Process Reward Model (PRM) evaluating trajectories in progress; and for medical imaging as a spatial attention signal concentrating capacity on diagnostically relevant regions. Empirically, a conditional diffusion framework achieves −17.3 dB mean normalized mean square error (NMSE) on 3GPP Urban Microcell (UMi) benchmarks; a verifier-guided adaptive test-time compute framework improves Llama-3.1-8B accuracy from 43.8% to 65.4% on MATH-500; and ZoomFormer achieves F1 = 0.899 on HLung-6 with only 4.23M parameters, outperforming substantially larger baselines.