Retrieval-Augmented Interpretable Learning: Towards Task-Specific Zero-Shot Models in Healthcare

2026-07-20Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors present RAIL, a new method that creates understandable prediction models for medical tasks without needing lots of new data. It uses descriptions of new tasks and past learned models to make predictions and explain them in simple features doctors understand. RAIL also shows how confident it is, so uncertain cases can get extra review instead of automatic decisions. This makes it useful for healthcare, where new and rare tasks appear often and clear, trustworthy predictions are needed. Their tests show RAIL works well even with very few examples and helps make the model’s decisions more transparent.

probabilistic meta-learningzero-shot learninginterpretable modelstask-specific predictorscoefficient-spaceclinical procedure predictionuncertainty quantificationfew-shot learningretrieval-augmented learningmodel transparency
Authors
Sazan Mahbub, Caleb Ellington, Zhiyuan Li, Yixin Yang, Souvik Kundu, Ben Lengerich, Eric P. Xing
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.