cs.LGJan 19, 2025
SaveModel Predictive Task Sampling for Efficient and Robust Adaptation
Organizations: Department of Automation, Tsinghua University · Informatics Institute, University of Amsterdam
Abstract
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.
Figures & tables
Figure 1: Framework of MPTS in Adaptation Learning. a. The left shows standard random sampling for generating candidate tasks. The middle in blue denotes costly evaluation of tasks (e.g., agent–environment interaction or foundation model forward pass) in DRM to select Top- worst ones. The middle in green depicts MPTS, which predicts task difficulty via a lightweight generative model, avoiding expensive evaluation. The right illustrates the standard optimization pipeline in Meta-Learning, DR, or SFT (Snow: frozen models; Fire: updated models). b. MPTS samples candidate tasks, ranks their difficulty via the risk predictive model’s predictions for subset selection, and updates the learner, approximating ’s Monte Carlo optimization in a predictive manner. The gathered risk signals further update RPMs online. c. The RPM utilizes the risk history to train under a streaming VI framework. d. The RPM simulates adaptation outcomes for candidate identifiers, computes acquisition scores, and selects the Top- identifiers for the -th iteration.
Figure 2: Fundamental Concepts: Task Identifiers, Episodic Learning and Probabilistic Graphical Models. a. The task distribution is uniform and defined over meaningful identifiers . For example, the amplitude and the phase specifies a sinusoid curve to complete with K-shot observed data points. Robots like Half-Cheetahs are trained to accomplish different locomotion tasks with varying masses and velocities. Some multimodal pattern recognition tasks’ identifiers are implicit but can be described from a reference model, e.g., text encoders in CLIP 67 . b. The tail task generalization corresponds to , i.e., the integral of tail task risk values in red. In OOD generalization, this work prompt-tunes CLIP on ImageNet 77 to test on ImageNet-S 92 . c. Here, the generative model includes grey units as observed variables and white ones as unobservable. The solid directed lines describe the generative model 90 . We use the dash-directed lines to indicate the recognition model and approximate inference within autoencoding variational Bayes 48 .
Figure 3: K-shot Sinusoid Regression Results (10 Runs). Note that a - i are results with MAML as the backbone, while j - n reports the results with Reptile as the backbone. a. Shown are curves of averaged MSEs with standard error of the mean (SEM) on the validation task set during meta-training for all methods. b. Curves illustrate the MSEs with SEM on the validation task set during the meta-training process. c. The meta-trained machine learners are tested on a fixed task set, reporting the average MSEs and CVaR values with SEM. d. Displayed are meta-testing results with SEM using MPTS machine learners trained by various ratios. e. Displayed are meta-testing results using MPTS machine learners trained in various pseudo batch sizes, i.e., . f. The PCC values with SEM are tracked during meta-training. g. At a specific iteration, the statistical correlation between predicted and exact adaptation risk values of the task batch is visualized with overall . h. The required relative run-time is computed for all methods during meta-training with ERM as the anchor. i. At some meta-training time-step, we visualize the subset selection from the pseudo batch under the RPM. j. We illustrate the curves of averaged MSEs with SEM on the validation task set during Reptile meta-training for all relevant baselines. k. We track the corresponding MSEs with SEM on the validation task set throughout the Reptile meta-training. l. Reported are the tested average MSEs and CVaR values with SEM using the meta-trained machine learners on a fixed task set. m. The PCC values with SEM are tracked for MPTS during Reptile meta-training. n. We compute the relative run-time for all methods during meta-training with ERM as the anchor.
| Dataset | Metrics | ERM | DRM | GDRM | OHTM | DATS | TDPS | EMTL | MPTS (Ours) | |
| ImageNet-CG 39 | 76.87 0.31 | 77.47 0.24 | 77.56 0.19 | 76.52 0.14 | 77.55 0.21 | 77.28 0.17 | 77.59 0.25 | 78.16 0.25 | +1.29 | |
| 82.04 0.21 | 82.47 0.20 | 82.62 0.20 | 81.34 0.17 | 82.31 0.18 | 82.33 0.24 | 82.56 0.14 | 82.95 0.26 | +0.91 | ||
| 84.62 0.24 | 84.99 0.17 | 85.16 0.22 | 84.06 0.29 | 85.19 0.14 | 85.29 0.23 | 85.14 0.17 | 85.33 0.26 | +0.71 | ||
| Avg | 89.22 0.22 | 89.54 0.16 | 89.49 0.18 | 88.59 0.17 | 89.45 0.12 | 89.69 0.11 | 89.34 0.21 | 89.86 0.14 | +0.64 | |
| ImageNet-CI 39 | 80.40 0.18 | 80.49 0.22 | 80.26 0.19 | 77.08 0.18 | 78.86 0.15 | 78.07 0.19 | 80.41 0.16 | 80.94 0.19 | +0.54 | |
| 84.98 0.17 | 85.05 0.19 | 84.90 0.22 | 81.60 0.11 | 83.66 0.22 | 83.16 0.11 | 85.04 0.18 | 85.70 0.24 | +0.72 |
Table 1: 5-Way 1-Shot Meta-Testing Classification Results on Various Datasets and Efficiency Comparison (10 Runs). We report testing , , and average accuracies with corresponding SEMs evaluated by the meta-trained machine learner on ID and OOD datasets. The best results are in bold with the runner-up underlined, and MPTS’s performance gains over ERM are marked in blue. With experiments on ImageNet-A as an example, we report the memory cost and runtime relative to ERM during meta-training.
Figure 4: Meta-RL Results on Mujoco Environments (10 Runs). a. The cumulative returns with SEM belonging to validation MDPs are displayed during meta-training. b. We compute the average cumulative returns with SEMs on validation MDPs during meta-training. c. Tracked are the RPM’s PCC values with SEMs over training iterations. d. The relative clock time quantifies the computational complexity for all methods on Walker2dVel, where ERM’s runtime works as the anchor. e. We report returns with SEM of meta-testing MDPs. f. The box-plot reports results averaged over meta-testing MDPs. g. With PEARL 70 as the Meta-RL backbone, we illustrate the learning curves and meta-testing results with SEM on HalfCheetahBody and HalfCheetahMass from RoML 33 baseline.
Figure 5: DR Results on Ergo-Reacher and Lunar-Lander (10 Runs). a. In Ergo-Reacher, the , , , and average cumulative returns with SEM on validation MDPs are reported. The RPM’s PCC values with SEM are tracked during DR training. b. In Lunar-Lander, the cumulative returns on validation MDPs with SEM are illustrated. The RPM’s PCC values with SEM are tracked during DR training. c. We test the DR-trained policies on the fixed MDP set and report the cumulative returns with SEM. d. The returns averaged over DR-testing MDPs are illustrated. e. The required runtime is computed for all methods on Lunar-Lander. f. In Lunar-Lander, shown are frequencies of sampled identifiers using MPTS during DR training. g. In Lunar-Lander, we test the trained policies in both in-distribution (ID) domains and out-of-distribution (OOD) domains to report each task’s average returns with SEM.
| Dataset | Metrics | ERM | DRM | GDRM | OHTM | MPTS (Ours) | |
| ImageNet 77 | 31.70 0.12 | 32.23 0.15 | 31.37 0.18 | 31.53 0.21 | 32.52 0.16 | +0.82 | |
| (ID) | 42.87 0.14 | 44.10 0.14 | 42.97 0.18 | 43.28 0.18 | 44.28 0.14 | +1.41 | |
| 51.43 0.10 | 52.61 0.11 | 51.75 0.19 | 51.96 0.16 | 52.72 0.15 | +1.29 | ||
| Avg | 70.80 0.08 | 70.90 0.10 | 71.00 0.08 | 70.80 0.10 | 71.20 0.09 | +0.40 | |
| ImageNet-A 40 | 15.35 0.21 | 15.50 0.25 | 15.62 0.22 | 15.44 0.25 | 18.38 0.22 | +3.03 | |
| (OOD) | 22.82 0.23 | 23.05 0.21 | 23.13 0.19 | 23.00 0.23 | 24.05 0.20 | +1.23 |
Table 2: Testing Classification Results after Prompt-Tuning on ImageNet and Efficiency Comparison (10 Runs). We report testing , , and average accuracies together with corresponding SEMs evaluated by the prompt-tuned machine learner on ID and OOD datasets. Evaluation on OOD datasets corresponds to the domain generalization setting. The best results are in bold with the runner-up underlined, and MPTS’s performance gains over ERM are marked in blue. During prompt-tuning ImageNet, we report the memory cost and runtime relative to ERM for all methods.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure S1 : Risk Predictive Model as the World Model to Guide Conditional Interaction. This shows the explanation of MPTS in the presence of active inference theory.
Figure S2 : Risk Predictive Module in MPTS for Active Subset Selection. MPTS adopts a predict-then-optimize strategy and uses a predictive module in green to approximately score the task subset difficulty and obtain the preferred task subset. While the traditional method in blue exhausts tasks in construction and evaluation to filter preferred subset.
Figure S3 : Comparison between Traditional Multi-Task Learning and Meta Learning and Domain Randomization. Traditional multi-task learning aims to exploit the correlation between heterogeneous tasks, and task IDs in a batch do not change 36 . In comparison, meta learning and domain randomization typically consider the homogeneous task family and resample task IDs in a batch to optimize. The latter is our studied scenario, referred to as task episodic optimization. [Same types of colors mean the homogeneous tasks with depth to denote the various episode.]
| ERM | DRM | GDRM | OHTM | MPTS (Ours) | |
|---|---|---|---|---|---|
| computation | |||||
| task eval |
Table S1 : Computational Complexities using Different Methods. Here, we drop out the ranking or reweighting computational complexity as the model complexity of the machine learner considered in this analysis is major, such as the multimodal foundation models. refers to the required iteration steps until the convergence for separate methods.
| Benchmarks | Identifier Meaning | Identifier Range |
| K-shot sinusoid regression | amplitude and phase | |
| Meta-RL: HalfCheetahMassVel | mass and velocity | |
| Meta-RL: HalfCheetahVel | velocity | |
| Meta-RL: ReacherPos | goal location | |
| Meta-RL: Walker2dMassVel | mass and velocity | |
| Meta-RL: Walker2dVel | velocity |
Table S2 : Benchmarks with Explicit Task Identifiers. Here, we list the detail information about the task identifier to induce the task distribution.
Figure S4 : The Encoder-Decoder Neural Network to Paramterize the RPM.
Figure S5 : Rank Flipping after One-Step Update. We number tasks from the sampled pseudo task batch for the -th subset selection. The left figure partitions the set into as illustrated in the Proof of Lemma 8 , with the blue ones, i.e., 1, 2, 3 as Top- tasks evaluated under and the red ones as the complementary set. Note that tasks 2 and 8 are swapped in relative Top- ranking after one-step model update; hence, the right figure depicts the Top- tasks as 1,3, and 8. The gray lines indicate possible order-flipped cross-pairs between and . The blue and the red lines are probably flipped cross-pairs in separate partitioned sets’ induced graphs.
| Benchmarks | Adaptation | Backbone | Expensive Part |
|---|---|---|---|
| K-shot sinusoid regression | few-shot | MAML/Reptile | computations |
| N-way K-shot image classification | few-shot | MaPLe | computation/memory |
| Meta-RL | few-shot | MAML/PEARL | interactions |
| Robotic DR | zero-shot | TD3 | interactions |
| SFT | many-shot | MaPLe | computation/memory |
Table S3 : A Summary of the Considered Benchmarks. Here, we list the primary expensive part in task episodic learning for each scenario together with backbone methods. Also note that N-way K-shot image classification and SFT requires implicit task identifiers while others can directly access explicit task identifiers as the lower dimensional task representation.
More in this topic
Adaptive Sampling
Adaptive sampling allocates limited data, measurements, or computation where they are expected to be most useful.
See all 202 papers