Paper deep dive
Retrieval-Augmented Interpretable Learning: Towards Task-Specific Zero-Shot Models in Healthcare
Sazan Mahbub, Caleb Ellington, Zhiyuan Li, Yixin Yang, Souvik Kundu, Ben Lengerich, Eric P. Xing
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 7/21/2026, 5:07:55 AM
Summary
The paper introduces Retrieval-Augmented Interpretable Learning (RAIL), a probabilistic meta-learning framework designed for zero-shot and few-shot generation of task-specific interpretable clinical prediction models. RAIL synthesizes coefficient-space structure by retrieving related source tasks from a parametric memory and transferring their learned coefficients to a new target task based on natural-language descriptions. The framework provides uncertainty estimates for retrieval, coefficients, and predictions, enabling reliability-aware deployment in healthcare settings where data is long-tailed and scarce. RAIL achieves 73.4% accuracy in zero-shot settings and 73.2% in extreme few-shot regimes, outperforming supervised baselines that degrade with limited data.
Entities (17)
Relation Signals (12)
RAIL → achievesaccuracy → 73.4%
confidence 98% · it achieves 73.4% accuracy in the held-out zero-shot settings
RAIL → achievesaccuracy → 73.2%
confidence 98% · remains near 73.2% accuracy in the extreme few-shot regime with only 2-4 examples
RAIL → provides → Uncertainty Estimation
confidence 95% · Because RAIL is probabilistic, it also produces uncertainty over retrieval, coefficients, and predictions.
RAIL → supports → Few-shot Learning
confidence 95% · When limited target-task labels are available, the same formulation supports few-shot posterior adaptation
RAIL → supports → Zero-shot Learning
confidence 95% · RAIL is primarily designed for zero-shot model synthesis
RAIL → uses → Parametric Memory
confidence 95% · RAIL assumes access to a memory of prior tasks, where each task is represented by both a language embedding and the coefficients of an interpretable linear model.
RAIL → evaluatedon → MIMIC-IV
confidence 92% · We evaluate RAIL on long-tailed clinical procedure prediction tasks... We construct clinical procedure prediction tasks from MIMIC-IV
Ben Lengerich → affiliatedwith → University of Wisconsin–Madison
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:We introduce Retrieval-Augmented Interpretable Learning (RAIL), a probabilistic meta-learning framework for zero-shot generation of task-specific interpretable models that synthesizes coefficient-space structure from natural-language task descriptions and a memory of previously learned task-specific predictors. RAIL retrieves related source tasks, transfers structure through coefficient space, and generates a new predictor in the original diagnostic-feature space, enabling zero-shot and few-shot clinical procedure prediction with feature-level explanations. Its probabilistic formulation provides uncertainty over retrieval, model coefficients, and predictions, supporting reliability-aware deployment: uncertain predictions or unstable explanations can be flagged for additional clinical review rather than treated as automatic decisions. This makes RAIL particularly suited for healthcare settings, where prediction tasks are highly long-tailed, new clinical targets arise frequently, and models must remain inspectable, uncertainty-aware, and compatible with human oversight. Across long-tailed clinical procedure prediction tasks, RAIL maintains reliable performance across data-availability regimes: it achieves 73.4% accuracy in the held-out zero-shot settings, where no supervised task-specific model can be trained, and remains near 73.2% accuracy in the extreme few-shot regime with only 2-4 examples, where supervised task-specific models perform close to chance. RAIL further benefits from clinically informed task representations and yields retrieval, uncertainty, and coefficient-level diagnostics that make model behavior more transparent. These results suggest a path toward scalable clinical prediction systems that can adapt to new tasks while preserving interpretability and reliability.
Tags
Links
- Source: https://arxiv.org/abs/2607.17508v1
- Canonical: https://arxiv.org/abs/2607.17508v1
Trouble viewing inline? Open PDF directly →
Full Text
65,053 characters extracted from source content.
Expand or collapse full text
Retrieval-Augmented Interpretable Learning: Towards Task-Specific Zero-Shot Models in Healthcare Sazan Mahbubc Caleb Ellingtonc,g Zhiyuan Liw Yixin Yangw Souvik Kundui Ben Lengerichw Eric P. Xingc,m,g cCarnegie Mellon University wUniversity of Wisconsin–Madison mMohamed bin Zayed University of AI gGenBio AI iIntel Abstract We introduce Retrieval-Augmented Interpretable Learning (RAIL), a probabilistic meta-learning framework for zero-shot generation of task-specific interpretable models that synthesizes coefficient-space structure from natural-language task descriptions and a memory of previously learned task-specific predictors. RAIL retrieves related source tasks, transfers structure through coefficient space, and generates a new predictor in the original diagnostic-feature space, enabling zero-shot and few-shot clinical procedure prediction with feature-level explanations. Its probabilistic formulation provides uncertainty over retrieval, model coefficients, and predictions, supporting reliability-aware deployment: uncertain predictions or unstable explanations can be flagged for additional clinical review rather than treated as automatic decisions. This makes RAIL particularly suited for healthcare settings, where prediction tasks are highly long-tailed, new clinical targets arise frequently, and models must remain inspectable, uncertainty-aware, and compatible with human oversight. Across long-tailed clinical procedure prediction tasks, RAIL maintains reliable performance across data-availability regimes: it achieves 73.4% accuracy in the held-out zero-shot settings, where no supervised task-specific model can be trained, and remains near 73.2% accuracy in the extreme few-shot regime with only 2–4 examples, where supervised task-specific models perform close to chance. RAIL further benefits from clinically informed task representations and yields retrieval, uncertainty, and coefficient-level diagnostics that make model behavior more transparent. These results suggest a path toward scalable clinical prediction systems that can adapt to new tasks while preserving interpretability and reliability. †footnotetext: Correspondence to: smahbub@cs.cmu.edu, lengerich@wisc.edu, epxing@cs.cmu.edu. 1 Introduction Recent advances in clinical predictive modeling have highlighted the need for models that are not only accurate, but also adaptive, interpretable, and reliable under limited supervision. In many healthcare settings, one needs task-specific predictors for deciding whether different procedures, medications, or interventions should be administered to a patient, conditioned on a shared set of diagnostic measurements. However, clinical task distributions are naturally long-tailed: common procedures may have abundant labels, while many clinically meaningful or newly considered tasks have only a few examples, and some have no task-specific labels at all. This makes independently training a supervised model for every task unreliable in precisely the regimes where adaptable clinical prediction is most needed. A growing body of work has studied meta-models and task-conditioned predictors that dynamically generate or adapt models for specific contexts or tasks (Hospedales et al., 2021; Lengerich et al., 2023; Ellington et al., 2025; Deuschel et al., 2024). In parallel, pretrained language models have been explored as sources of prior knowledge for downstream decision-making (Thirunavukarasu et al., 2023), including reinforcement learning (Du et al., 2023; Karimpanal et al., 2023; Zhang and Lu, 2024), feature selection (Adila et al., 2024), causal discovery (Long et al., 2023; Liu et al., 2024), and healthcare retrieval-augmented prediction pipelines (Jin et al., 2024). These approaches suggest that task descriptions and pretrained representations can provide useful priors for adaptation. Yet, many language-model-based or black-box transfer systems do not directly produce transparent, task-specific predictive models. Their outputs often remain embedded in latent representations or natural-language rationales, limiting their use when clinicians need inspectable feature-level weights, uncertainty estimates, and signals for when additional human review is warranted. This motivates our central question: Can we generate an interpretable task-specific clinical prediction model in a zero-shot manner, using only a natural-language task description and a memory of previously learned interpretable predictors? Such a framework would combine the adaptability of meta-learning and retrieval-augmented modeling with the transparency of classical interpretable predictors. It would also support clinical oversight by exposing which features drive a prediction, when the generated model is uncertain, and when predictions should be flagged for additional review. We introduce Retrieval-Augmented Interpretable Learning (RAIL), a probabilistic meta-learning framework for zero-shot generation of task-specific interpretable models, synthesizing coefficient-space structure from natural-language task descriptions and a memory of previously learned linear predictors. RAIL represents each prior task using both a language embedding of its task description and the coefficients of an interpretable model trained for that task. Given a new target task, RAIL retrieves related source tasks, transfers structure through coefficient space, and generates a new predictor in the original diagnostic-feature space. The resulting model is directly inspectable: each coefficient corresponds to an input feature, enabling feature-level interpretation rather than post-hoc explanation of an opaque predictor. The key idea is that task descriptions and task-specific coefficients encode complementary forms of transferable knowledge. Language embeddings provide semantic alignment across procedures and clinical targets, while learned coefficients encode how diagnostic features were used by prior interpretable predictors. RAIL combines these signals through a probabilistic retrieval formulation: retrieved tasks define a retrieval-conditioned prior over target-task coefficients, which directly yields a generated predictor in the zero-shot regime and can be further adapted when limited target-task labels are available. Because RAIL is probabilistic, it also produces uncertainty over retrieval, coefficients, and predictions. These uncertainty signals can identify ambiguous retrieval neighborhoods, unstable feature-level explanations, and individual predictions that may require clinician review. We evaluate RAIL on long-tailed clinical procedure prediction tasks, where each task asks whether a procedure should be administered to a patient given diagnostic results. Our experiments show that RAIL is especially useful in low-data and zero-shot regimes, where task-specific supervision is scarce or unavailable. Quantitatively, RAIL remains stable as supervision decreases: it obtains 73.4% accuracy on held-out zero-shot tasks using only task descriptions and retrieved prior predictors, and achieves 73.2% accuracy in the extreme few-shot regime with only 2–4 labeled examples. In contrast, supervised task-specific baselines degrade substantially in this regime, highlighting the value of transferring coefficient-space structure from related source tasks. Controlled baselines and ablations show that performance depends on clinically informed task representations, meaningful retrieval, and learned coefficient synthesis. Retrieval diagnostics further show that RAIL retrieves coherent and oracle-consistent task neighborhoods, while embedding perturbation studies confirm that task-representation quality is central to model generation. Finally, uncertainty analyses show that predictive confidence tracks empirical accuracy, uncertainty identifies likely failures, selective prediction reduces risk, and coefficient-level uncertainty helps distinguish stable from unstable explanations. Figure 1: Overview of RAIL. Retrieved task embeddings and learned linear models from the parametric memory condition the generation of a task-specific interpretable predictor. Our contributions are as follows: • We introduce RAIL, a probabilistic meta-learning framework for zero-shot generation of task-specific interpretable clinical prediction models from natural-language task descriptions and a memory of prior interpretable predictors. • We formulate model generation as retrieval-conditioned coefficient-space transfer, producing predictors in the original diagnostic-feature space and preserving direct feature-level interpretability. • We develop a probabilistic formulation that supports zero-shot generation, optional few-shot adaptation, and uncertainty estimation over retrieval, coefficients, and predictions. • We evaluate RAIL on long-tailed clinical procedure prediction tasks, with baselines, retrieval ablations, text-encoder comparisons, oracle retrieval diagnostics, and embedding perturbation studies showing when and why retrieval-augmented generation succeeds. • We demonstrate that RAIL provides reliability and interpretability diagnostics, including uncertainty-based failure detection, selective prediction, and coefficient-level uncertainty for identifying stable versus unstable explanations. 2 RAIL: Retrieval-Augmented Interpretable Learning Goal. Given a new task t described only by natural language ctc_t, our goal is to synthesize an interpretable task-specific predictor ft(x)=σ(θt⊤x)f_t(x)=σ( _t x) without training a new black-box model from scratch. RAIL assumes access to a memory of prior tasks, where each task is represented by both a language embedding and the coefficients of an interpretable linear model. VItI_tθt _tctc_ttD_t Figure 2: PGM for RAIL. Observed variables are shaded. The embeddings provide semantic alignment across tasks, while the coefficients provide transferable parameter-level structure in the original feature space. RAIL is primarily designed for zero-shot model synthesis: at inference time, no labeled examples from task t are required. When limited target-task labels are available, the same formulation supports few-shot posterior adaptation by combining the retrieval-conditioned prior with a task-specific estimator. 2.1 Parametric memory: semantic keys and interpretable values We assume training tasks i=(i,ci)i=1n\T_i=(D_i,c_i)\_i=1^n, where i=(xij,yij)j=1miD_i=\(x_ij,y_ij)\_j=1^m_i share a common feature space ⊂ℝfX ^f, and each task has a natural-language description cic_i. For each task i, we train an interpretable linear model, e.g., logistic regression, to obtain coefficients θ^i∈ℝf θ_i ^f. We also compute a task embedding Ei=LM(ci)∈ℝeE_i=LM(c_i) ^e using a frozen language-model encoder. The resulting parametric memory is =(θ^i,Ei)i=1n, =\( θ_i,E_i)\_i=1^n, (1) where embeddings serve as semantic retrieval keys and coefficients serve as interpretable parameter-space values. 2.2 Latent-variable formulation RAIL introduces two latent quantities. The retrieval latent It⊆I_t is a retrieved subset of memory entries selected from the full parametric memory, with |It|=S|I_t|=S. Each element of ItI_t is a source-task memory entry (θ^j,Ej)( θ_j,E_j). The task-parameter latent θt∈ℝf _t ^f is the coefficient vector of the interpretable predictor for task t. 2.3 Generative model The parametric memory V and task description ctc_t are observed conditioning variables. RAIL models uncertainty over both the retrieved memory subset and the task-specific coefficients: pϕ(t,θt,It∣ct,)=p(t∣θt)pϕ(θt∣It,ct)p(It∣ct,). p_φ(D_t, _t,I_t c_t,V)=p(D_t _t)\;p_φ( _t I_t,c_t)\;p(I_t c_t,V). (2) The overview and the probabilistic graphical model are shown in Fig. 1 and Fig. 2, respectively. Likelihood. For binary prediction, we use the logistic GLM likelihood p(y∣x,θ)=Bernoulli(σ(θ⊤x)), p(y x,θ)=Bernoulli (σ(θ x) ), (3) where (x,y)∈t(x,y) _t. For regression, the likelihood can analogously be written as p(y∣x,θ)=(y;θ⊤x,σ2)p(y x,θ)=N(y;θ x,σ^2). Retrieval prior. The prior p(It∣ct,)p(I_t c_t,V) is induced by cosine similarity between the target task embedding Et=LM(ct)E_t=LM(c_t) and memory entries. Let st,jpriors^prior_t,j denote the cosine-similarity score for memory entry j. In implementation, we retrieve a candidate subset of size S using these scores and define normalized prior weights within this candidate set: πt,jprior=exp(st,jprior/τp)∑ℓ=1Sexp(st,ℓprior/τp). π^prior_t,j= (s^prior_t,j/ _p) _ =1^S (s^prior_t, / _p). (4) This cosine-based prior provides a semantic anchor for the learned retriever. Retrieval-conditioned coefficient prior. For each retrieved memory entry (θ^j,Ej)∈It( θ_j,E_j)∈ I_t, define Kj=(θ^j∥Ej),Vj=θ^j. K_j=( θ_j\,\|\,E_j), V_j= θ_j. (5) Conditioned on ItI_t and ctc_t, a multi-head cross-attention generator aggregates the retrieved coefficient values and outputs a Gaussian prior over the target-task coefficients: pϕ(θt∣It,ct)=(θt;μϕ,t,diag(σϕ,t2)), p_φ( _t I_t,c_t)=N\! ( _t; _φ,t,diag(σ^2_φ,t) ), (6) where (μϕ,t,σϕ,t2)( _φ,t,σ^2_φ,t) are produced by neural heads applied to the attention output (Vaswani et al., 2017). We use a diagonal covariance and parameterize logσϕ,t2 σ^2_φ,t. 2.4 Learned retrieval and coefficient inference We use the factorized variational family qψ(It,θt∣t,ct,)=qψ(It∣ct,)qψ(θt∣t,It,ct). q_ψ(I_t, _t _t,c_t,V)=q_ψ(I_t c_t,V)\;q_ψ( _t _t,I_t,c_t). (7) The retrieval factor qψ(It∣ct,)q_ψ(I_t c_t,V) is restricted to be label-free, so that the learned retriever can be used in zero-shot inference. It is trained jointly with the coefficient generator to improve over the cosine-similarity prior, while remaining semantically anchored to it through the retrieval KL term in Eq. (15). The coefficient factor qψ(θt∣t,It,ct)q_ψ( _t _t,I_t,c_t) is the variational posterior used when target-task labels are available. Amortized learned retriever. For each candidate memory entry, we compute a learned retrieval score from the target embedding and candidate key: st,jpost=gψ(Et,Kj). s^post_t,j=g_ψ(E_t,K_j). (8) Let tK_t denote the top-k candidates under these scores, with k=|t|<Sk=|K_t|<S. To approximate near-discrete retrieval while preserving differentiability, we sharpen only positive top-k scores: s~t,jpost=st,jpost/τq,j∈t and st,jpost>0,st,jpost,otherwise,0<τq≤1. s^post_t,j= casess^post_t,j/ _q,&j _t and s^post_t,j>0,\\ s^post_t,j,&otherwise, cases 0< _q≤ 1. (9) We then define the learned retrieval distribution qψ(It∣ct,)=Cat(πtpost),πt,jpost=exp(s~t,jpost)∑ℓ=1Sexp(s~t,ℓpost). q_ψ(I_t c_t,V)=Cat\! (π^post_t ), π^post_t,j= ( s^post_t,j) _ =1^S ( s^post_t, ). (10) Although ItI_t denotes the retrieved memory subset, the implementation uses a tractable categorical relaxation over the S candidate entries in the retrieved set. During training, the soft weights πtpostπ^post_t are used as differentiable retrieval weights inside the attention generator. During inference, retrieval is performed using qψ(It∣ct,)q_ψ(I_t c_t,V) rather than the cosine-similarity prior. Coefficient posterior. We use a Gaussian posterior with diagonal covariance: qψ(θt∣t,It,ct)=(θt;μq,t,diag(σq,t2)). q_ψ( _t _t,I_t,c_t)=N\! ( _t; _q,t,diag(σ^2_q,t) ). (11) For stability, our implementation shares posterior and prior variances, σq,t2≡σϕ,t2σ^2_q,t≡σ^2_φ,t, and constructs the posterior mean as μq,t=α⊙μϕ,t+(1−α)⊙θ^tself, _q,t=α _φ,t+(1-α) θ^self_t, (12) where θ^tself θ^self_t is an optional task-specific estimator, such as a logistic-regression model trained on tD_t, and α∈(0,1)fα∈(0,1)^f is a learned per-coordinate gate. Thus, the posterior interpolates between the retrieval-conditioned prior and the task-specific estimator, relying more on the prior when target-task supervision is scarce. Reparameterization and prediction. We sample θt=μq,t+σq,t⊙ε,ε∼(0,I), _t= _q,t+ _q,t , (0,I), (13) and predict with the GLM likelihood, e.g., σ(θt⊤x)σ( _t x) for binary classification. 2.5 Variational objective with learned retrieval The task evidence marginalizes over both the retrieved memory index and the target-task coefficients: pϕ(t∣ct,)=∑It∫pϕ(t,θt,It∣ct,)θt. p_φ(D_t c_t,V)= _I_t p_φ(D_t, _t,I_t c_t,V)\,d _t. (14) Using the restricted variational family in Eq. (7), we optimize the negative ELBO: Training objective. t(ϕ,ψ) _t(φ,ψ) =−qψ(It∣ct,)qψ(θt∣t,It,ct)[logp(t∣θt)] =-E_q_ψ(I_t c_t,V)E_q_ψ( _t _t,I_t,c_t) [ p(D_t _t) ] (15) +qψ(It∣ct,)[KL(qψ(θt∣t,It,ct)∥pϕ(θt∣It,ct))] +E_q_ψ(I_t c_t,V) [KL\! (q_ψ( _t _t,I_t,c_t)\, \|\,p_φ( _t I_t,c_t) ) ] +KL(qψ(It∣ct,)∥p(It∣ct,)). +KL\! (q_ψ(I_t c_t,V)\, \|\,p(I_t c_t,V) ). The likelihood term trains generated coefficients to explain target-task labels, the coefficient KL anchors the adapted posterior to the retrieval-conditioned prior, and the retrieval KL regularizes the learned retriever toward the semantic similarity prior. Thus, downstream prediction encourages retrieval of memory entries useful for coefficient synthesis, while the KL terms prevent unconstrained drift. A full derivation and additional theoretical analyses are given in Appendix F. 3 Results and Discussion We evaluate RAIL on clinical procedure prediction tasks, where each task predicts whether a procedure should be administered to a patient given diagnostic measurements. This setting is naturally long-tailed: common procedures have abundant supervision, while many clinically meaningful procedures have few examples or are evaluated zero-shot. Our experiments assess: (i) low-data and zero-shot predictive performance, (i) the role of meaningful retrieval, (i) whether task embeddings recover transferable parameter neighborhoods, and (iv) whether RAIL’s uncertainty estimates support reliability and interpretability. 3.1 Experimental Setup Dataset. We construct clinical procedure prediction tasks from MIMIC-IV (Johnson et al., 2023). Each task predicts whether a procedure is administered from 217217 diagnostic laboratory features, using the procedure definition as the task description. The dataset contains 7,4877,487 procedure tasks: 4,4124,412 with at least two positive examples for supervised/few-shot evaluation and 3,0753,075 held-out tasks for zero-shot evaluation. Dataset construction, normalization, and task-regime details are given in Appendix B. Tasks and regimes. Each procedure defines a binary prediction task t with diagnostic features x∈ℝfx ^f (where f=217f=217), a natural-language task description ctc_t, and task-specific labels. Since all tasks share the same diagnostic-feature space, task-specific linear models can be represented in a common coefficient space. We evaluate across sample-scarcity regimes with 50+50+, 1010–4949, 55–99, and 22–44 examples, and a held-out zero-shot regime. In zero-shot evaluation, target tasks are completely excluded from the retrieval memory: their descriptions, patient samples, labels, and task-specific coefficients are unseen during inference. RAIL receives only ctc_t, retrieves source tasks distinct from the target task, and deploys the retrieval-conditioned prior mean μϕ,t _φ,t as the synthesized model. Baselines, ablations, and metrics. We compare against supervised task-specific references, direct retrieval-transfer baselines, retrieval perturbations, and task-text encoder variants. Supervised references use target-task labels and are therefore not zero-shot methods. Retrieval-transfer baselines include Top-11 retrieved LR and Top-k coefficient averaging. Retrieval ablations replace normal retrieval with random or semantically distant source tasks. For task-text representations, we compare DistilBERT-base-uncased with the larger medical-domain MedEmbed-large-v0.1. We report task-averaged accuracy as the primary predictive metric, and use retrieval diagnostics, uncertainty-based failure detection, selective prediction, and coefficient-level uncertainty to evaluate interpretability and reliability. 3.2 Long-Tailed Task Distribution Motivates Retrieval-Augmented Model Generation Clinical procedure tasks are highly imbalanced in sample availability: a small number of procedures have abundant supervision, while many have only a few labeled examples. Figure 3: Long-tailed task distribution. Procedure tasks are highly imbalanced in sample availability. This long-tail structure makes direct task-specific training unreliable in the low-data tail and motivates retrieval-augmented model generation, where related prior tasks provide parameter-level structure for new or data-scarce tasks. Fig. 3 shows the long-tailed task distribution. Fig. 4 and Appendix Tab. 3 compare RAIL with the task-specific logistic regression across sample regimes. Here, ‘LR oracle’ denotes an LR model fit separately for each target task. In high-resource tasks, the supervised reference remains competitive, as expected. As target-task supervision decreases, however, the LR degrades sharply, dropping from approximately 0.760.76 to 0.420.42 in F1 and from 0.760.76 to 0.550.55 in accuracy between the 50+50+ and 22–44 regimes. In contrast, RAIL remains comparatively stable, with F1 between 0.700.70 and 0.750.75 and accuracy between 0.720.72 and 0.750.75, indicating that retrieval-conditioned model generation provides useful prior structure when direct supervision is scarce. 3.3 Baselines and Main Ablations Figure 4: RAIL remains stable across sample-scarcity regimes, while the supervised task-specific LR oracle degrades in the low-data tail. We compare RAIL with supervised references, a retrieval-free meta-learning baseline, retrieval perturbations, and direct retrieval-transfer baselines. In addition to accuracy, Tab. 1 reports whether each method exposes clinically useful diagnostics: an inspectable predictor, task-text conditioning, retrieval traceability, and coefficient uncertainty. Table 1: Baseline and main ablation comparison. Diagnostic columns indicate whether a method provides an inspectable predictor, uses task-description embeddings, exposes retrieved source-task traces, or estimates coefficient uncertainty. For repeated runs, we report mean ± standard deviation over five independent runs. Variant Inspectable Task text Retrieval trace Coeff. uncertainty Few-shot acc. Zero-shot acc. Supervised references Task-specific MLP × × × × 0.5330.533 N/A Task-specific XGBoost × × × × 0.5000.500 N/A Task-specific Random Forest × × × × 0.5730.573 N/A Task-specific Logistic Regression ✓ × × × 0.5500.550 N/A Task-specific Decision Tree ✓ × × × 0.5310.531 N/A RAIL variants Full RAIL (w/ MedEmbed) ✓ ✓ ✓ ✓ 0.732±0.0090.732± 0.009 0.734±0.0110.734± 0.011 Full RAIL (w/ DistilBERT) ✓ ✓ ✓ ✓ 0.721±0.0100.721± 0.010 0.718±0.0130.718± 0.013 Without retrieval (w/ MedEmbed) ✓ ✓ × ✓ 0.703±0.0120.703± 0.012 0.701±0.0140.701± 0.014 Retrieval perturbation ablations RAIL + random retrieval ✓ ✓ ✓ ✓ 0.657±0.0190.657± 0.019 0.659±0.0210.659± 0.021 RAIL + distant retrieval ✓ ✓ ✓ ✓ 0.579±0.0080.579± 0.008 0.587±0.0100.587± 0.010 Retrieval-transfer baselines Top-11 retrieved LR ✓ ✓ ✓ × 0.5870.587 0.5850.585 Top-k coefficient average (k=100k=100) ✓ ✓ ✓ × 0.6560.656 0.6520.652 Tab. 1 shows that full RAIL provides the strongest combination of predictive performance and clinical inspectability. MedEmbed improves over DistilBERT, suggesting that clinically informed task-text representations provide stronger priors for procedure prediction. The retrieval-free variant is competitive, empirically proving that task descriptions carry useful prior; however, full RAIL improves over it by grounding model generation in retrieved source-task coefficients and exposing a retrieval trace. Random and distant retrieval substantially degrade performance, showing that retrieval must be semantically meaningful. Direct retrieval-transfer baselines also underperform full RAIL, indicating that learned coefficient synthesis is more effective than direct model reuse or naive coefficient averaging. Thus, RAIL’s benefit is not only higher low-data and zero-shot accuracy, but also the combination of feature-level inspectability, retrieval-grounded traceability, and uncertainty-aware diagnostics. 3.4 Stress-Testing Task-Embedding Space Beyond the main predictive ablations in Tab. 1, we stress-test the task-representation space through perturbations, compression, and randomization. Additional robustness results are provided in Appendix Fig. 10. These analyses show that Gaussian noise degrades few-shot and zero-shot performance, compact embeddings preserve useful retrieval structure, and random or shuffled embeddings collapse retrieval quality. Together, these results indicate that RAIL depends on semantic organization in the task-embedding space, while useful retrieval structure can still be compressed. 3.5 Retrieval and Embedding Diagnostics (a) Query vs. retrieved neighbors. (b) Pairwise similarity among neighbors. Figure 5: Embedding–model alignment. Top-k retrieved tasks are closer to the query and more coherent among themselves than random neighbors. RAIL assumes that semantically related task descriptions retrieve source tasks with transferable coefficient structure. We test this through embedding–model alignment and oracle retrieval consistency. Embedding–model alignment. Figure 6: Oracle retrieval consistency. Embedding retrieval better matches oracle-relevant tasks than random retrieval and TF-IDF. We first ask whether nearest neighbors in task-embedding space are also close in coefficient space. Fig. 5 compares top-k retrieved neighbors with random neighbors. Across embedding variants, retrieved neighbors are substantially more similar to the query task than random neighbors, and retrieved neighborhoods are more internally coherent. Similarity decreases as k grows, as broader neighborhoods include less related tasks. This supports RAIL’s central mechanism: local neighborhoods in task-embedding space contain transferable parameter information. Oracle retrieval consistency. We further compare embedding retrieval with TF-IDF and random retrieval using an oracle-overlap criterion. As shown in Fig. 6, embedding retrieval achieves higher oracle overlap across retrieval depths than random retrieval and is competitive with or better than TF-IDF. Thus, task embeddings capture similarity relevant for model transfer rather than only lexical overlap. 3.6 Uncertainty Diagnostics and Selective Prediction (a) Confidence vs. accuracy. (b) Error vs. uncertainty. Figure 7: Predictive uncertainty diagnostics. (a) Predictive confidence tracks empirical accuracy. (b) Higher predictive uncertainty corresponds to higher error rates. RAIL provides retrieval, coefficient, and predictive uncertainty. Figure 8: Selective prediction. Lower-uncertainty predictions yield lower risk at reduced coverage. We focus on predictive uncertainty for reliability analysis: we sample K=100K=100 coefficient vectors from the posterior, compute the induced predictive probabilities, and use the binary entropy of the posterior-mean probability as the uncertainty score. Prediction confidence is max(p¯,1−p¯) ( p,1- p). Prediction confidence aligns with empirical correctness (Fig. 7(a) and 7(b)): accuracy increases monotonically from low- to high-confidence bins, while error rate increases from low- to high-uncertainty bins. We further define a sample-level failure as [y^≠y]1[ y≠ y] and evaluate uncertainty scores using AUROC and AUPRC. Predictive uncertainty is the strongest failure-detection signal, achieving AUROC 0.6830.683 and AUPRC 0.4190.419, outperforming retrieval entropy and parameter uncertainty; ROC and precision–recall curves are provided in Appendix Fig. 11. Thus, output-level uncertainty is most aligned with prediction-level failure, while retrieval and coefficient uncertainty provide complementary task- and model-level diagnostics. Finally, we evaluate uncertainty-based abstention by ranking examples by increasing uncertainty and retaining the lowest-uncertainty subset. Fig. 8 shows that risk decreases as coverage decreases, suggesting that uncertain predictions can be flagged for expert review rather than used automatically. 3.7 Coefficient-Level Interpretability with Uncertainty RAIL generates linear models in the original diagnostic-feature space, so each coefficient corresponds to an inspectable input feature. Its posterior over coefficients further enables uncertainty-aware interpretation. Figure 9: Feature-level interpretability with uncertainty. RAIL ranks diagnostic features by posterior coefficient magnitude and reports uncertainty intervals. For each task and feature, we record posterior coefficient mean, magnitude, variance, and standard deviation (Fig. 9). Stable features have large posterior importance with narrow intervals, while uncertain features have wider intervals and should be interpreted cautiously. Additional stable and unstable coefficient examples are provided in Appendix Fig. 12. This coefficient-level uncertainty supports feature-level inspection while making explanation stability explicit. 3.8 Summary and Limitations Overall, RAIL is most useful in the low-data tail of clinical procedure prediction, where task-specific supervised models degrade and zero-shot synthesis becomes necessary. Retrieval and embedding diagnostics show that its gains depend on meaningful task representations and semantically aligned source-task retrieval, while ablations show that learned coefficient synthesis improves over direct transfer or naive averaging. Uncertainty analyses further show that RAIL can flag unreliable predictions and unstable explanations for review. Limitations are that RAIL assumes a shared diagnostic-feature space and depends on task-memory coverage; when no related source tasks exist, retrieval-conditioned synthesis may be less reliable. Coefficient explanations are also associational rather than causal. 4 Conclusions We introduced RAIL, a retrieval-augmented probabilistic meta-learning framework for generating task-specific interpretable clinical prediction models from natural-language task descriptions. By retrieving prior task predictors and synthesizing structure directly in coefficient space, RAIL produces models in the original diagnostic-feature space, preserving feature-level interpretability while supporting uncertainty-aware prediction. Across long-tailed clinical procedure tasks, RAIL is most effective when target-task supervision is scarce or unavailable, and our ablations show that its gains rely on clinically informed task representations, meaningful retrieval, and learned coefficient synthesis. Together, these results suggest a practical path toward scalable clinical prediction systems that can adapt to newly emerging or low-resource tasks without sacrificing inspectability, uncertainty awareness, or compatibility with human oversight. Acknowledgments and Disclosure of Funding This research has been graciously funded by the National Science Foundation (NSF) awards BCS2040381, CNS2414087 and IIS2123952 (to S.M. and E.X.); the Defense Advanced Research Projects Agency (DARPA) award HR00112390063 (to S.M. and E.X.); the Semiconductor Research Corporation (SRC) AIHW award 2024AH3210 (to S.M. and E.X.); and, the National Institutes of Health (NIH) award R01GM140467 (to C.E. and E.X.). Any opinions, findings, and conclusions or recommendations expressed in this publication are those of the author(s) and do not necessarily reflect the views of the National Science Foundation, the Defense Advanced Research Projects Agency, the Semiconductor Research Corporation, and the National Institutes of Health. References [1] D. Adila, C. Shin, L. Cai, and F. Sala (2024) Zero-shot robustification of zero-shot models. External Links: 2309.04344, Link Cited by: Appendix A, §1. [2] S. Bordt, B. Lengerich, H. Nori, and R. Caruana (2024) Data science with llms and interpretable models. External Links: 2402.14474, Link Cited by: Appendix A. [3] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch (2021) Decision transformer: reinforcement learning via sequence modeling. External Links: 2106.01345, Link Cited by: Appendix A. [4] K. Choi, C. Cundy, S. Srivastava, and S. Ermon (2022) LMPriors: pre-trained language models as task-specific priors. External Links: 2210.12530, Link Cited by: Appendix A. [5] J. Deuschel, C. Ellington, Y. Luo, B. Lengerich, P. Friederich, and E. P. Xing (2024) Contextualized policy recovery: modeling and interpreting medical decisions with adaptive imitation learning. In Forty-first International Conference on Machine Learning, Cited by: Appendix A, §1. [6] Y. Du, O. Watkins, Z. Wang, C. Colas, T. Darrell, P. Abbeel, A. Gupta, and J. Andreas (2023) Guiding pretraining in reinforcement learning with large language models. External Links: 2302.06692, Link Cited by: Appendix A, §1. [7] C. N. Ellington, B. J. Lengerich, T. B. Watkins, J. Yang, A. K. Adduri, S. Mahbub, H. Xiao, M. Kellis, and E. P. Xing (2025) Learning to estimate sample-specific transcriptional networks for 7,000 tumors. Proceedings of the National Academy of Sciences 122 (21), p. e2411930122. Cited by: Appendix A, §1. [8] Y. Gao, Y. Xiong, X. Gao, K. Jia, J. Pan, Y. Bi, Y. Dai, J. Sun, and H. Wang (2023) Retrieval-augmented generation for large language models: a survey. arXiv preprint arXiv:2312.10997. Cited by: Appendix A. [9] T. Hospedales, A. Antoniou, P. Micaelli, and A. Storkey (2021) Meta-learning in neural networks: a survey. IEEE transactions on pattern analysis and machine intelligence 44 (9), p. 5149–5169. Cited by: Appendix A, §1. [10] M. Jin, Q. Yu, D. Shu, C. Zhang, L. Fan, W. Hua, S. Zhu, Y. Meng, Z. Wang, M. Du, et al. (2024) Health-llm: personalized retrieval-augmented disease prediction system. arXiv preprint arXiv:2402.00746. Cited by: Appendix A, §1. [11] A. E. Johnson, L. Bulgarelli, L. Shen, A. Gayles, A. Shammout, S. Horng, T. J. Pollard, S. Hao, B. Moody, B. Gow, et al. (2023) MIMIC-iv, a freely accessible electronic health record dataset. Scientific data 10 (1), p. 1. Cited by: Appendix B, §3.1. [12] T. G. Karimpanal, L. B. Semage, S. Rana, H. Le, T. Tran, S. Gupta, and S. Venkatesh (2023) LaGR-seq: language-guided reinforcement learning with sample-efficient querying. External Links: 2308.13542, Link Cited by: Appendix A, §1. [13] B. Lengerich, B. Aragam, and E. P. Xing (2019) Learning sample-specific models with low-rank personalized regression. Advances in Neural Information Processing Systems 32. Cited by: Appendix A. [14] B. Lengerich, C. N. Ellington, A. Rubbi, M. Kellis, and E. P. Xing (2023) Contextualized machine learning. arXiv preprint arXiv:2310.11340. Cited by: Appendix A, §1. [15] P. Lewis, E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal, H. Küttler, M. Lewis, W. Yih, T. Rocktäschel, et al. (2020) Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems 33, p. 9459–9474. Cited by: Appendix A. [16] S. Li, X. Puig, C. Paxton, Y. Du, C. Wang, L. Fan, T. Chen, D. Huang, E. Akyürek, A. Anandkumar, et al. (2022) Pre-trained language models for interactive decision-making. Advances in Neural Information Processing Systems 35, p. 31199–31212. Cited by: Appendix A. [17] C. Liu, Y. Chen, T. Liu, M. Gong, J. Cheng, B. Han, and K. Zhang (2024) Discovery of the hidden world with large language models. External Links: 2402.03941, Link Cited by: Appendix A, §1. [18] S. Long, A. Piché, V. Zantedeschi, T. Schuster, and A. Drouin (2023) Causal discovery with language models as imperfect experts. External Links: 2307.02390, Link Cited by: Appendix A, §1. [19] M. Mahbub, S. Srinivasan, E. Begoli, and G. D. Peterson (2022) BioADAPT-mrc: adversarial learning-based domain adaptation improves biomedical machine reading comprehension task. Bioinformatics 38 (18), p. 4369–4379. Cited by: Appendix A. [20] S. Mahbub, S. Kundu, and E. P. Xing (2025) PRISM: enhancing protein inverse folding through fine-grained retrieval on structure-sequence multimodal representations. arXiv preprint arXiv:2510.11750. Cited by: Appendix A. [21] H. A. MD (2025) Meta-learning approaches for causal discovery in dynamic healthcare and robotics environments. Mesopotamian Journal of Artificial Intelligence in Healthcare 2025, p. 136–153. Cited by: Appendix A. [22] Y. Tan, C. Yang, X. Wei, C. Chen, W. Liu, L. Li, J. Zhou, and X. Zheng (2022) Metacare++: meta-learning with hierarchical subtyping for cold-start diagnosis prediction in healthcare data. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, p. 449–459. Cited by: Appendix A. [23] A. J. Thirunavukarasu, D. S. J. Ting, K. Elangovan, L. Gutierrez, T. F. Tan, and D. S. W. Ting (2023) Large language models in medicine. Nature medicine 29 (8), p. 1930–1940. Cited by: Appendix A, §1. [24] S. Tonekaboni, S. Joshi, M. D. McCradden, and A. Goldenberg (2019) What clinicians want: contextualizing explainable machine learning for clinical end use. In Machine learning for healthcare conference, p. 359–380. Cited by: Appendix A. [25] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin (2017) Attention is all you need. Advances in neural information processing systems 30. Cited by: §2.3. [26] X. Yan, Y. Song, X. Feng, M. Yang, H. Zhang, H. B. Ammar, and J. Wang (2025) Efficient reinforcement learning with large language model priors. In 13th International Conference on Learning Representations Iclr 2025, p. 30818–30842. Cited by: Appendix A. [27] W. Zhang and Z. Lu (2024) AdaRefiner: refining decisions of language models with adaptive feedback. External Links: 2309.17176, Link Cited by: Appendix A, §1. [28] X. S. Zhang, F. Tang, H. H. Dodge, J. Zhou, and F. Wang (2019) Metapred: meta-learning for clinical risk prediction with limited patient electronic health records. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining, p. 2487–2495. Cited by: Appendix A. Appendix A Related Works Clinical predictive modeling often faces a data-scarcity problem: many clinically relevant tasks have limited supervision, yet the resulting models must remain accurate, interpretable, and reliable enough to support expert review. Several research directions address different parts of this challenge. Contextualized machine learning studies prediction settings where the input–output relationship varies across contexts, individuals, or tasks, which is common in clinical and decision-making domains [14, 5, 24, 13, 7]. Meta-learning similarly seeks to use experience across tasks to support rapid adaptation to new tasks with limited data [21, 28, 22, 9]. Recent meta-models and task-conditioned predictors extend this idea by dynamically generating or adapting predictors for particular contexts or tasks [9, 14, 7, 5]. These methods provide a foundation for scalable low-data prediction, but many rely on latent representations or flexible neural predictors rather than directly inspectable models in the original feature space. Another approach is to use retrieval to provide task-relevant external information. Retrieval-augmented generation conditions a model on retrieved content, allowing it to use information beyond what is stored in its parameters [15, 8]. In biomedical domain, retrieval-augmented pipelines have been used to incorporate health reports and medical knowledge into large models for improved feature extraction and prediction [10, 20]. Retrieval is therefore a natural mechanism for grounding predictions in relevant prior information. However, standard RAG systems typically retrieve documents or knowledge snippets for language generation or latent prediction, rather than using retrieved prior predictors as reusable task-level structure. A complementary direction uses pretrained language models as sources of prior knowledge when labeled data are scarce [23, 19]. Language-model priors have been used in reinforcement learning to guide exploration and decision refinement [3, 6, 12, 27], in feature selection to identify useful variables [1], and in causal discovery or reasoning to propose or evaluate plausible causal structure [18, kıcıman2024causalreasoninglargelanguage, 17]. Related work has also studied language-model priors across multiple decision-making settings [4, 16, 26]. These studies show that pretrained representations can encode useful task-level information, but this information is often used to guide black-box decision processes, construct latent features, or generate natural-language rationales rather than directly producing task-specific interpretable predictors. Finally, work combining language models with interpretable models aims to leverage language-derived knowledge while preserving transparency in the final decision process [2]. This direction is especially relevant for healthcare, where users may need to inspect feature-level evidence, assess reliability, and decide when additional review is needed. Overall, prior work offers important tools for low-data prediction—task adaptation, retrieval, language-derived priors, and interpretability—but leaves open the challenge of combining these ingredients into a framework that produces task-specific predictors that are simultaneously data-efficient, directly inspectable, and uncertainty-aware. Appendix B Dataset and Preprocessing We use the MIMIC-IV dataset [11] to construct clinical procedure prediction tasks. The preprocessed data consist of two main tables: laboratory events and procedure records. Both tables are grouped by hospital admission and patient identifiers, (hadm_id,subject_id)( hadm\_id, subject\_id). Laboratory events provide diagnostic measurements, which form the input features, while procedure records define task-specific binary labels indicating whether a procedure was administered during a hospital stay. We retain only admissions with corresponding laboratory measurements so that every labeled example has an associated diagnostic feature vector. Each input vector contains 217217 diagnostic features derived from laboratory measurements. Since laboratory items have different clinical reference ranges, we normalize each raw measurement using its item-specific lower and upper reference values: normalized_feature=raw_value−lower_rangeupper_range−lower_range.normalized\_feature= raw\_value-lower\_rangeupper\_range-lower\_range. (16) Values outside the reference interval are not clipped; this preserves clinically meaningful deviations while placing heterogeneous diagnostics on a comparable scale. Each clinical procedure defines a separate binary prediction task. The dataset contains 7,4877,487 procedure tasks in total. Among these, 4,4124,412 procedures have at least two positive examples and are used for high-resource, low-resource, and few-shot evaluation. For each of these tasks, we construct balanced task-specific datasets with equal numbers of positive and negative examples and use a 50/5050/50 train–test split within the task. The remaining 3,0753,075 procedure tasks are used for zero-shot evaluation, where no target-task training examples or task-specific oracle models are available during inference. Tab. 2 summarizes the resulting task regimes. We use natural-language procedure definitions as task descriptions. Example descriptions include “Respiratory Ventilation, 24–96 Consecutive Hours,” “Inspection of Upper Intestinal Tract, Via Natural or Artificial Opening Endoscopic,” “Fluoroscopy of Multiple Coronary Arteries using Other Contrast,” and “Introduction of Other Antineoplastic into Central Vein, Percutaneous Approach.” These descriptions provide the task-text signal used by RAIL for retrieval-augmented model generation. Table 2: Summary of task regimes used in RAIL evaluation. Each clinical procedure defines one binary prediction task. The zero-shot regime contains held-out target procedures whose labels and task-specific oracle models are excluded from the retrieval memory during inference. Regime Task ID range Examples per task High-resource 0–341341 50+50+ Moderate-resource 342342–15211521 1010–4949 Low-resource 15221522–23062306 55–99 Few-shot 23072307–44114411 22–44 Zero-shot 44124412–74867486 No target-task labels used Appendix C Hardware Specifications All experiments were run on a single workstation equipped with 12 CPU cores of AMD EPYC 9354 processor at 3.25 GHz, 100 GB of RAM, and one NVIDIA RTX A6000 GPU with 48 GB of VRAM. Appendix D Additional Experimental Details This section provides supplementary tables for the experimental setup and analyses in the main paper. Tab. 3 reports the full performance breakdown across sample-scarcity regimes. Tab. 4 summarizes the ablation and diagnostic suite. Tab. 5 defines the uncertainty signals used in reliability analyses, and Tab. 6 summarizes the interpretable outputs exposed by RAIL. Table 3: Performance across sample-scarcity regimes. The task-specific LR oracle uses target-task labels and remains strongest when sufficient supervision is available, while RAIL is most beneficial in low-data regimes. Sample regime RAIL LR oracle Precision Recall F1 Accuracy Precision Recall F1 Accuracy 50+50+ 0.7033 0.7883 0.7396 0.7286 0.7639 0.7641 0.7610 0.7615 1010–4949 0.7258 0.7788 0.7396 0.7343 0.6953 0.6689 0.6597 0.6715 55–99 0.7455 0.8002 0.7503 0.7441 0.5975 0.5745 0.5437 0.6016 22–44 (few-shot) 0.6878 0.7876 0.7129 0.7293 0.3860 0.5285 0.4240 0.5500 Zero-shot 0.6621 0.7837 0.7027 0.7348 N/A N/A N/A N/A Table 4: Summary of RAIL ablation and diagnostic analyses. Analysis group Variants tested Purpose Text encoder DistilBERT, MedEmbed Tests effect of clinical task representations. Retrieval perturbation Normal, random, distant Tests causal role of meaningful retrieval. Transfer baseline Top-11 LR, Top-k coefficient average Tests learned synthesis vs direct transfer. Embedding corruption Gaussian noise Tests sensitivity to semantic corruption. Embedding compression PCA-32, PCA-64 Tests whether retrieval structure is compressible. Embedding randomization Random, shuffled Tests whether structure exceeds chance. Retrieval diagnostics Query-neighbor, neighbor-neighbor, Oracle Recall@k Tests semantic/model alignment. Uncertainty diagnostics Retrieval, parameter, predictive uncertainty Tests reliability and failure detection. Table 5: Uncertainty signals used in RAIL diagnostics. Signal Definition Level Retrieval uncertainty −∑jqψ(It=j)logqψ(It=j)- _jq_ψ(I_t=j) q_ψ(I_t=j) Task Parameter uncertainty ∑rσq,t,r2 _rσ^2_q,t,r Task/model Predictive uncertainty −p¯logp¯−(1−p¯)log(1−p¯)- p p-(1- p) (1- p) Sample Confidence max(p¯,1−p¯) ( p,1- p) Sample Table 6: Interpretability outputs produced by RAIL. Output Level Interpretation Generated coefficients θt _t Task Interpretable diagnostic-feature weights. Posterior coefficient variance Feature Uncertainty in each feature’s contribution. Top-k feature frequency Feature Stability of feature importance across posterior samples. Retrieval posterior weights Task memory Attribution to retrieved source tasks. Prediction confidence Sample Reliability of individual predictions. Appendix E Additional Retrieval Robustness, Uncertainty, and Interpretability Diagnostics This section provides additional diagnostic figures supporting the retrieval robustness, uncertainty, and coefficient-level interpretability analyses in the main paper. Fig. 10 reports controlled perturbations of the task-representation space. Fig. 11 reports ROC and precision–recall curves for uncertainty-based failure detection. Fig. 12 shows examples of stable and unstable coefficient-level explanations. (a) Noise perturbation. (b) Embedding dimension. (c) Retrieval ablations. Figure 10: Additional representation ablations. Noise degrades RAIL performance, compact embeddings preserve useful retrieval structure, and random or shuffled embeddings collapse retrieval quality. (a) Failure-detection ROC. (b) Failure-detection PR. Figure 11: Failure detection using uncertainty. Predictive uncertainty is the strongest signal for identifying likely incorrect predictions, outperforming retrieval entropy and parameter uncertainty in both ROC and precision–recall analyses. Figure 12: Stable and unstable coefficient examples. Stable tasks exhibit top coefficients with relatively narrow uncertainty intervals, while unstable tasks show wider intervals for some features, indicating explanations that should be interpreted with caution. Appendix F Additional Theoretical Analysis F.1 ELBO derivation for learned retrieval We derive the objective used in Eq. (15). Recall the conditional generative model pϕ(t,θt,It∣ct,)=p(t∣θt)pϕ(θt∣It,ct)p(It∣ct,). p_φ(D_t, _t,I_t c_t,V)=p(D_t _t)\;p_φ( _t I_t,c_t)\;p(I_t c_t,V). (17) Here, ItI_t denotes the retrieved memory object induced by the retriever. In the exact set-valued view, It⊆I_t . In implementation, we use a tractable top-S categorical relaxation over the retrieved candidate entries; the closed-form categorical KL below corresponds to this implemented relaxation. The marginal evidence for task t is pϕ(t∣ct,)=∑It∫pϕ(t,θt,It∣ct,)θt, p_φ(D_t c_t,V)= _I_t p_φ(D_t, _t,I_t c_t,V)\,d _t, (18) where the summation is over the support of the retrieval distribution. For any variational distribution qψ(θt,It∣t,ct,)q_ψ( _t,I_t _t,c_t,V), we can write logpϕ(t∣ct,) p_φ(D_t c_t,V) =log∑It∫qψ(θt,It∣t,ct,)pϕ(t,θt,It∣ct,)qψ(θt,It∣t,ct,)θt = _I_t q_ψ( _t,I_t _t,c_t,V) p_φ(D_t, _t,I_t c_t,V)q_ψ( _t,I_t _t,c_t,V)\,d _t (19) ≥qψ(θt,It∣t,ct,)[logpϕ(t,θt,It∣ct,)qψ(θt,It∣t,ct,)], _q_ψ( _t,I_t _t,c_t,V) [ p_φ(D_t, _t,I_t c_t,V)q_ψ( _t,I_t _t,c_t,V) ], (20) where the inequality follows from Jensen’s inequality. Thus, ℒELBO(t)=qψ[logpϕ(t,θt,It∣ct,)−logqψ(θt,It∣t,ct,)]. _ELBO^(t)=E_q_ψ [ p_φ(D_t, _t,I_t c_t,V)- q_ψ( _t,I_t _t,c_t,V) ]. (21) We use the restricted amortized factorization qψ(θt,It∣t,ct,)=qψ(It∣ct,)qψ(θt∣t,It,ct). q_ψ( _t,I_t _t,c_t,V)=q_ψ(I_t c_t,V)\;q_ψ( _t _t,I_t,c_t). (22) The retrieval factor is intentionally label-free. This restriction preserves zero-shot retrieval at inference time, where tD_t is unavailable. It is still a valid variational family, but it may yield a looser bound than a fully label-conditioned retrieval posterior. Substituting Eq. (17) and Eq. (22) into Eq. (21) gives ℒELBO(t) _ELBO^(t) =qψ[logp(t∣θt)+logpϕ(θt∣It,ct)+logp(It∣ct,) =E_q_ψ [ p(D_t _t)+ p_φ( _t I_t,c_t)+ p(I_t c_t,V) (23) −logqψ(θt∣t,It,ct)−logqψ(It∣ct,)]. - q_ψ( _t _t,I_t,c_t)- q_ψ(I_t c_t,V) ]. Grouping the likelihood, coefficient, and retrieval terms yields ℒELBO(t) _ELBO^(t) =qψ(It∣ct,)qψ(θt∣t,It,ct)[logp(t∣θt)] =E_q_ψ(I_t c_t,V)E_q_ψ( _t _t,I_t,c_t) [ p(D_t _t) ] (24) +qψ(It∣ct,)qψ(θt∣t,It,ct)[logpϕ(θt∣It,ct)qψ(θt∣t,It,ct)] +E_q_ψ(I_t c_t,V)E_q_ψ( _t _t,I_t,c_t) [ p_φ( _t I_t,c_t)q_ψ( _t _t,I_t,c_t) ] +qψ(It∣ct,)[logp(It∣ct,)qψ(It∣ct,)]. +E_q_ψ(I_t c_t,V) [ p(I_t c_t,V)q_ψ(I_t c_t,V) ]. Recognizing the last two terms as negative KL divergences gives ℒELBO(t) _ELBO^(t) =qψ(It∣ct,)qψ(θt∣t,It,ct)[logp(t∣θt)] =E_q_ψ(I_t c_t,V)E_q_ψ( _t _t,I_t,c_t) [ p(D_t _t) ] (25) −qψ(It∣ct,)[KL(qψ(θt∣t,It,ct)∥pϕ(θt∣It,ct))] -E_q_ψ(I_t c_t,V) [KL\! (q_ψ( _t _t,I_t,c_t)\, \|\,p_φ( _t I_t,c_t) ) ] −KL(qψ(It∣ct,)∥p(It∣ct,)). -KL\! (q_ψ(I_t c_t,V)\, \|\,p(I_t c_t,V) ). This leads to the negative ELBO objective in Eq. (15). The first term trains the generated coefficients to explain the target-task labels. The second term anchors the coefficient posterior to the retrieval-conditioned coefficient prior. The third term trains the learned retriever relative to the cosine-similarity retrieval prior. Closed-form KL terms. When both the coefficient posterior and prior are diagonal Gaussians, qψ(θt∣⋅)=(θt;μq,t,diag(σq,t2)),pϕ(θt∣⋅)=(θt;μϕ,t,diag(σϕ,t2)), q_ψ( _t ·)=N( _t; _q,t,diag(σ^2_q,t)), p_φ( _t ·)=N( _t; _φ,t,diag(σ^2_φ,t)), (26) the general diagonal Gaussian KL is KL(qψ(θt∣⋅)∥pϕ(θt∣⋅))=12∑r=1f[logσϕ,t2(r)σq,t2(r)+σq,t2(r)+(μq,t(r)−μϕ,t(r))2σϕ,t2(r)−1]. (q_ψ( _t ·)\,\|\,p_φ( _t ·) )= 12 _r=1^f [ _φ,t^2(r) _q,t^2(r)+ _q,t^2(r)+( _q,t^(r)- _φ,t^(r))^2 _φ,t^2(r)-1 ]. (27) In our implementation, we tie the posterior and prior variances, σq,t2≡σϕ,t2σ^2_q,t≡σ^2_φ,t, so Eq. (27) simplifies to KL(qψ(θt∣⋅)∥pϕ(θt∣⋅))=12∑r=1f(μq,t(r)−μϕ,t(r))2σϕ,t2(r). (q_ψ( _t ·)\,\|\,p_φ( _t ·) )= 12 _r=1^f ( _q,t^(r)- _φ,t^(r))^2 _φ,t^2(r). (28) For the implemented categorical retrieval relaxation, qψ(It∣⋅)=Cat(πtpost)q_ψ(I_t ·)=Cat( _t^post) and p(It∣⋅)=Cat(πtprior)p(I_t ·)=Cat( _t^prior), the retrieval KL is KL(qψ(It∣⋅)∥p(It∣⋅))=∑j=1Sπt,jpostlogπt,jpostπt,jprior. (q_ψ(I_t ·)\,\|\,p(I_t ·) )= _j=1^Sπ^post_t,j π^post_t,jπ^prior_t,j. (29) Practical estimator. We estimate the expected log-likelihood term with one reparameterized coefficient sample θt _t per minibatch and minimize a weighted negative-ELBO objective: ^t=BCE(y,σ(θt⊤x))⏟data term+βθKLθ⏟(28)+βIKLI⏟(29). J_t= BCE (y,σ( _t x) )_data term+ _θ KL_θ_ eq:kl-theta+ _I KL_I_ eq:kl-I. (30) The weights βθ _θ and βI _I control the strength of coefficient-prior anchoring and retrieval-prior anchoring, respectively. When βθ=βI=1 _θ= _I=1, this corresponds to the negative ELBO above; otherwise it is a weighted negative-ELBO objective. F.2 Prior anchoring and generalization We give a simple Rademacher-complexity argument showing why anchoring the coefficient posterior near a retrieval-conditioned prior can reduce the effective hypothesis class in low-data regimes. This result is not intended to characterize the full neural retrieval generator. Rather, it formalizes the intuition that, once the retrieved prior mean is fixed, constraining coefficients to remain near that mean controls the complexity of the induced linear predictors. Assumption 1 (Bounded features and bounded Lipschitz loss). The input features satisfy ‖x‖2≤R\|x\|_2≤ R almost surely. The loss ℓ(y^,y) ( y,y) is LℓL_ -Lipschitz in the prediction y y and takes values in an interval of length at most CℓC_ over the hypothesis class considered below. Proposition 1 (Generalization benefit of prior anchoring). Fix a retrieval-conditioned prior mean μϕ,t _φ,t independently of the target-task sample used in the bound, and consider the hypothesis class ℋB=x↦θ⊤x:‖θ−μϕ,t‖2≤B. _B=\x θ x:\|θ- _φ,t\|_2≤ B\. (31) Then the empirical Rademacher complexity satisfies ℜm(ℋB)≤RBm. R_m(H_B)≤ RB m. (32) Moreover, under Assumption 1, with probability at least 1−δ1-δ over a sample of size m, for all θ such that ‖θ−μϕ,t‖2≤B\|θ- _φ,t\|_2≤ B, ℒ(θ)≤ℒ^(θ)+2Lℓℜm(ℋB)+3Cℓlog(2/δ)2m≤ℒ^(θ)+2LℓRBm+3Cℓlog(2/δ)2m. (θ)≤ L(θ)+2L_ \, R_m(H_B)+3C_ (2/δ)2m≤ L(θ)+ 2L_ RB m+3C_ (2/δ)2m. (33) Proof. Let Sm=x1,…,xmS_m=\x_1,…,x_m\ be a fixed sample and let σ1,…,σm _1,…, _m be independent Rademacher variables. The empirical Rademacher complexity of ℋBH_B is ℜm(ℋB)=σ[sup‖θ−μϕ,t‖2≤B1m∑i=1mσiθ⊤xi]. R_m(H_B)=E_σ [ _\|θ- _φ,t\|_2≤ B 1m _i=1^m _iθ x_i ]. (34) Write θ=μϕ,t+uθ= _φ,t+u, where ‖u‖2≤B\|u\|_2≤ B. Then ℜm(ℋB) R_m(H_B) =σ[sup‖u‖2≤B1m∑i=1mσi(μϕ,t+u)⊤xi] =E_σ [ _\|u\|_2≤ B 1m _i=1^m _i( _φ,t+u) x_i ] (35) =σ[1m∑i=1mσiμϕ,t⊤xi+sup‖u‖2≤Bu⊤(1m∑i=1mσixi)]. =E_σ [ 1m _i=1^m _i _φ,t x_i+ _\|u\|_2≤ Bu ( 1m _i=1^m _ix_i ) ]. (36) The first term vanishes in expectation over the Rademacher variables because σ[σi]=0E_σ[ _i]=0. Therefore, ℜm(ℋB)=σ[sup‖u‖2≤Bu⊤(1m∑i=1mσixi)]. R_m(H_B)=E_σ [ _\|u\|_2≤ Bu ( 1m _i=1^m _ix_i ) ]. (37) By Cauchy–Schwarz, sup‖u‖2≤Bu⊤(1m∑i=1mσixi)≤Bm‖∑i=1mσixi‖2. _\|u\|_2≤ Bu ( 1m _i=1^m _ix_i )≤ Bm \| _i=1^m _ix_i \|_2. (38) Thus, ℜm(ℋB)≤Bmσ‖∑i=1mσixi‖2. R_m(H_B)≤ BmE_σ \| _i=1^m _ix_i \|_2. (39) Applying Jensen’s inequality, σ‖∑i=1mσixi‖2 _σ \| _i=1^m _ix_i \|_2 ≤(σ‖∑i=1mσixi‖22)1/2. ≤ (E_σ \| _i=1^m _ix_i \|_2^2 )^1/2. (40) Expanding the squared norm gives σ‖∑i=1mσixi‖22 _σ \| _i=1^m _ix_i \|_2^2 =σ[∑i=1m∑j=1mσiσjxi⊤xj] =E_σ [ _i=1^m _j=1^m _i _jx_i x_j ] (41) =∑i=1m‖xi‖22, = _i=1^m\|x_i\|_2^2, (42) because [σiσj]=0E[ _i _j]=0 for i≠ji≠ j and [σi2]=1E[ _i^2]=1. Since ‖xi‖2≤R\|x_i\|_2≤ R, ∑i=1m‖xi‖22≤mR2. _i=1^m\|x_i\|_2^2≤ mR^2. (43) Therefore, ℜm(ℋB)≤BmmR2=RBm. R_m(H_B)≤ Bm mR^2= RB m. (44) It remains to connect this complexity bound to generalization. By the standard Rademacher generalization theorem for a loss class with range at most CℓC_ , with probability at least 1−δ1-δ, uniformly for all h∈ℋBh _B, ℒ(h)≤ℒ^(h)+2ℜm(ℓ∘ℋB)+3Cℓlog(2/δ)2m. (h)≤ L(h)+2 R_m( _B)+3C_ (2/δ)2m. (45) Since ℓ is LℓL_ -Lipschitz in the prediction, the contraction inequality gives ℜm(ℓ∘ℋB)≤Lℓℜm(ℋB). R_m( _B)≤ L_ R_m(H_B). (46) Substituting the Rademacher-complexity bound derived above yields ℒ(θ)≤ℒ^(θ)+2Lℓℜm(ℋB)+3Cℓlog(2/δ)2m≤ℒ^(θ)+2LℓRBm+3Cℓlog(2/δ)2m. (θ)≤ L(θ)+2L_ R_m(H_B)+3C_ (2/δ)2m≤ L(θ)+ 2L_ RB m+3C_ (2/δ)2m. (47) This completes the proof. ∎ Appendix G Broader Impacts The paper discusses potential benefits for scalable, interpretable clinical prediction in low-data settings and emphasizes responsible use through uncertainty awareness, explanation stability, and clinician oversight. Appendix H Declaration of LLM usage The core method development in this research does not involve LLMs as any important, original, or non-standard components. LLMs were used solely for writing refinement and not for retrieval, discovery, or research ideation.