# Model Predictive Task Sampling for Efficient and Robust Adaptation

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

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 76 of its 113 references that are in the map.
- First posted on arXiv 19 Jan 2025.

## Builds on (most cited first, 30 of 76)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Flamingo: a Visual Language Model for Few-Shot Learning](https://mlsanity.com/paper/2204.14198/) (2022)
- [Prototypical Networks for Few-shot Learning](https://mlsanity.com/paper/1703.05175/) (2017)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [CLIPScore: A Reference-free Evaluation Metric for Image Captioning](https://mlsanity.com/paper/2104.08718/) (2021)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Learning to Prompt for Vision-Language Models](https://mlsanity.com/paper/2109.01134/) (2021)
- [The Many Faces of Robustness: A Critical Analysis of Out-of-Distribution Generalization](https://mlsanity.com/paper/2006.16241/) (2020)
- [Do ImageNet Classifiers Generalize to ImageNet?](https://mlsanity.com/paper/1902.10811/) (2019)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Conditional Prompt Learning for Vision-Language Models](https://mlsanity.com/paper/2203.05557/) (2022)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Natural Adversarial Examples](https://mlsanity.com/paper/1907.07174/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [HyperNetworks](https://mlsanity.com/paper/1609.09106/) (2016)
- [On First-Order Meta-Learning Algorithms](https://mlsanity.com/paper/1803.02999/) (2018)
- [Learning Robust Global Representations by Penalizing Local Predictive Power](https://mlsanity.com/paper/1905.13549/) (2019)
- [Meta-Learning in Neural Networks: A Survey](https://mlsanity.com/paper/2004.05439/) (2020)
- [Solving Rubik's Cube with a Robot Hand](https://mlsanity.com/paper/1910.07113/) (2019)
- [MaPLe: Multi-modal Prompt Learning](https://mlsanity.com/paper/2210.03117/) (2022)
- [Deep Learning on a Data Diet: Finding Important Examples Early in Training](https://mlsanity.com/paper/2107.07075/) (2021)
- [Beyond neural scaling laws: beating power law scaling via data pruning](https://mlsanity.com/paper/2206.14486/) (2022)
- [Meta-Learning with Implicit Gradients](https://mlsanity.com/paper/1909.04630/) (2019)

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

- [Can Prompt Difficulty be Online Predicted for Accelerating RL Finetuning of Reasoning Models?](https://mlsanity.com/paper/2507.04632/) (2025)
- [Prompt Curriculum Learning for Efficient LLM Post-Training](https://mlsanity.com/paper/2510.01135/) (2025)
- [FlyLoRA: Boosting Task Decoupling and Parameter Efficiency via Implicit Rank-Wise Mixture-of-Experts](https://mlsanity.com/paper/2510.08396/) (2025)
- [CurES: From Gradient Analysis to Efficient Curriculum Learning for Reasoning LLMs](https://mlsanity.com/paper/2510.01037/) (2025)
- [Utility-Diversity Aware Online Batch Selection for LLM Supervised Fine-tuning](https://mlsanity.com/paper/2510.16882/) (2025)
- [BOTS: A Unified Framework for Bayesian Online Task Selection in LLM Reinforcement Finetuning](https://mlsanity.com/paper/2510.26374/) (2025)
- [Small Generalizable Prompt Predictive Models Can Steer Efficient RL Post-Training of Large Reasoning Models](https://mlsanity.com/paper/2602.01970/) (2026)
- [Fast and Robust: Task Sampling with Posterior and Diversity Synergies for Adaptive Decision-Makers in Randomized Environments](https://mlsanity.com/paper/2504.19139/) (2025)
- [Fly-CL: A Fly-Inspired Framework for Enhancing Efficient Decorrelation and Reduced Training Time in Pre-trained Model-based Continual Representation Learning](https://mlsanity.com/paper/2510.16877/) (2025)
- [Dynamics-Predictive Sampling for Active RL Finetuning of Large Reasoning Models](https://mlsanity.com/paper/2603.10887/) (2026)
- [Selective Learning for Deep Time Series Forecasting](https://mlsanity.com/paper/2510.25207/) (2025)
- [Kuramoto Orientation Diffusion Models](https://mlsanity.com/paper/2509.15328/) (2025)
- [InspireDebate: Multi-Dimensional Subjective-Objective Evaluation-Guided Reasoning and Optimization for Debating](https://mlsanity.com/paper/2506.18102/) (2025)
- [Transformer-Based Spatial-Temporal Counterfactual Outcomes Estimation](https://mlsanity.com/paper/2506.21154/) (2025)
- [DTop-p MoE: Sparsity-Controlled Dynamic Top-p MoE for Foundation Model Pre-training](https://mlsanity.com/paper/2512.13996/) (2025)
- [Limited Reasoning Space: The cage of long-horizon reasoning in LLMs](https://mlsanity.com/paper/2602.19281/) (2026)
- [SpidR-Adapt: A Universal Speech Representation Model for Few-Shot Adaptation](https://mlsanity.com/paper/2512.21204/) (2025)
- [Gains: Fine-grained Federated Domain Adaptation in Open Set](https://mlsanity.com/paper/2510.15967/) (2025)
- [PERO: Efficient Robust Post-Training Foundation Models for Encrypted Traffic Classification](https://mlsanity.com/paper/2608.15504/) (2026)
- [You Only Edit Once: Incentivizing In-Context Capability of LLMs via Local Demonstration Refinement](https://mlsanity.com/paper/2609.33609/) (2026)
- [Bridging Risk Approximation Gaps in Model Predictive Task Sampling via In-Context Modeling](https://mlsanity.com/paper/2610.14913/) (2026)

## Abstract (arXiv)

Foundation models have revolutionized general-purpose problem-solving, offering rapid task adaptation through pretraining, meta-training, and finetuning. Recent crucial advances in these paradigms reveal the importance of challenging task prioritized sampling to enhance adaptation robustness under distribution shifts. However, ranking task difficulties over iteration as a preliminary step typically requires exhaustive task evaluation, which is practically unaffordable in computation and data-annotation. This study provides a novel perspective to illuminate the possibility of leveraging the dual importance of adaptation robustness and learning efficiency, particularly in scenarios where task evaluation is risky or costly, such as iterative agent-environment interactions for robotic policy evaluation or computationally intensive inference steps for finetuning foundation models. Firstly, we introduce Model Predictive Task Sampling (MPTS), a framework that bridges the task space and adaptation risk distributions, providing a theoretical foundation for robust active task sampling. MPTS employs a generative model to characterize the episodic optimization process and predicts task-specific adaptation risk via posterior inference. The resulting risk predictive model amortizes the costly evaluation of task adaptation performance and provably approximates task difficulty rankings. MPTS seamlessly integrates into zero-shot, few-shot, and supervised finetuning settings. Empirically, we conduct extensive experiments in pattern recognition using foundation models and sequential decision-making. Our results demonstrate that MPTS significantly enhances adaptation robustness for tail risk or out-of-distribution (OOD) tasks and improves learning efficiency compared to state-of-the-art (SoTA) methods. The code is available at the project site https://github.com/thu-rllab/MPTS.