AdaptMI: Adaptive Skill-based In-context Math Instructions for Small Language Models

We thank the reviewer for their positive comments on our work. Please find responses to your questions below.

“An explanation or analysis of the factors contributing to the instances, where AdaptMI+ method did not consistently outperform AdaptMI”

Response: We thank the reviewer for the question. As the only difference between AdaptMI+ and AdaptMI is the skill identification strategy (AdaptMI+ used an LLM to identify missing skills, while AdaptMI directly retrieved with skills from the Skill Bank), it is hard to statistically quantify when a different set of selected skill-based examples contributes to a better/worse performance.

To give a small example, we conducted a small-scale case study on the questions where the small model initially failed, got corrected with AdaptMI, but again failed with AdaptMI+. A common pattern observed in such cases is that AdaptMI+ tends to overfit to the small model’s previous mistake, rather than addressing the broader skill required to solve the question.

For instance, in this question (in MATH “Precalculus”):

Find the spherical coordinates of the point diametrically opposite $\\left( 3, \\frac{3 \\pi}{8}, \\frac{\\pi}{5} \\right).$

In this case, AdaptMI correctly identified the relevant skill as three_dimensional_geometry, highlighting the need to reason over all three spherical coordinates. In contrast, AdaptMI+ identified the missing skill as geometry_and_space_calculation, which narrowly targeted the model's previous mistake and failed to generalize to the full scope of the question. As a result, the small model failed again. We believe a deeper analysis on the differences between AdaptMI and AdaptMI+ can be an interesting future direction.

“Could the authors elaborate on the rationale for not including a comparison between methods AdaptMI and AdaptMI+ and the LLM?”

Response: Below we present the comparison between our approaches and the LLM performance (gpt-4o-mini). On GSM8K, the LLM is on par with Qwen2.5-7B-Instruct. While on MATH, the LLM is only on par with the 3B model, lacking behind the 7B model by ~7%.

We did not include this comparison in the earlier paper version because our approach only relied on minimal, high-level outputs (i.e., providing skill labels) from the LLM, instead of distilling its reasoning capabilities. The performance gain of AdaptMI mainly stems from its adaptive nature.

We thank the reviewer for their careful review and feedback. Please find our responses to the concerns raised in the review.

1. “High computational overhead in AdaptMI+ due to training an additional process reward model and using a large LLM”

We answer the question in two parts:

Re: Overhead due to training an additional process reward model (PRM)

In our experiments, the PRM was used to classify easy and difficult questions for the small model based on its responses. As mentioned in line 108, this was done to avoid access to ground truth labels. The PRM that we use is a general purpose model that is applicable to a wide range of math questions beyond MATH and GSM8K, even though it was trained on MATH and GSM8K questions. This indicates the potential to extend our method to new datasets without the need to train a specialized PRM for each one.

Table A: PRM performance for Qwen2.5-7B-Instruct

Motivated by the reviewer’s question, we also explored alternative approaches to classify questions as easy or difficult without relying on the PRM model. Specifically, we experimented with two heuristic strategies: (a) using the length of the model’s responses as a proxy and classifying questions with longer responses as difficult, and (b) measuring the consistency of the model across five sampled generations per question and classifying questions with lower consistency as difficult. Preliminary results suggest that both methods yield reasonably accurate predictions. However, we leave a more thorough investigation into the robustness and generalizability of these strategies in relation to PRM-based classification for future work.

Table: Performance of Qwen2.5-1.5B-Instruct on MATH

Re: Overhead of using an LLM

First, in AdaptMI, we use an LLM to label the skills required for each question and construct a corresponding skill bank. This labeling is performed once offline, introducing only a limited overhead. Despite this minimal reliance on the LLM, we demonstrate that it can lead to substantial improvements in the performance of the SLM.

Building on this, we proposed AdaptMI+, an enhanced strategy where the LLM is additionally used to annotate the missing skills in the model’s responses. While we acknowledge that this approach incurs greater overhead due to repeated LLM usage, it also highlights that such iterative access can provide meaningful benefits for model performance.

2. “Improvement is expected, …, but not due to a technical advancement, rather because it is indirectly leveraging the reasoning results of the larger model”

Our framework is inherently a teacher-student setup, where we want a stronger model to improve the performance of a weaker model. This setup will only become more prominent over time, as we build stronger AI models in the future and leverage them to improve the weaker models.

Secondly, we show that the teacher doesn’t need to give the ground truth answers directly to the student. Instead, we only utilize the metacognitive ability of the LLM to predict very high-level information about questions.

• Below is an example of the extracted LLM output:

algebraic_skills, understanding_circle_properties_and_algebraic_manipulation, coordinate_geometry_and_transformation_skills

The LLM is solely giving out skills without solving the problem or giving additional information.

• There are only 118 skills for MATH and 14 for GSM8K, indicating that these skills are high-level instructions rather than a strong guidance for problem solving.

More importantly, the teacher model (GPT-4o) is not necessarily stronger: it achieves only 69% on MATH, 5% lower than the baseline performance of Qwen2.5-7B-Instruct, and comparable to Qwen2.5-3B-Instruct. We will include this discussion in the next version.

We thank the reviewer for their extensive feedback on our work. Please find our responses to your questions below.

1. “Additional ablations on $\tau_2$ and its effect on the model’s overall performance”

Response: We thank the reviewer for the question. Below are ablation studies on the effect of varying $\tau_2$ or $\tau_1$ on small language models’ performance.

As presented in Table A, entirely removing $\tau_2$ (i.e., set $\tau_2=0$) will indeed increase the prediction accuracy and precision, yet it decreases the recall and F1, which are indicators of the detection rate of model failures. Therefore, in Table B below, the model performances under $\tau_2=0$ turn worse. We observe the same trend on removing $\tau_1$. We will include these observations in our paper revisions.

Table A: Reward model prediction metrics (accuracy / precision / recall / F1) across different thresholds for Qwen2.5-1.5B-Instruct on MATH.

Table B: Final performance of Qwen2.5-1.5B-Instruct on MATH under different thresholds for AdaptMI.

2. “Specific cases where consistency@5 does better/is on par with their approach and why to use one over the other”

Response: Thank you for the suggestion. We will add the following content in Section 3.2:

“While AdaptMI+ surpasses consistency@5 performance on most domains, it slightly lags behind on certain subjects such as Geometry and Precalculus for 1B or 3B models. These subjects are relatively difficult for the model, as suggested by their loss scores compared to other subjects (please check Table 9 in the appendix). Since AdaptMI requires models to have sufficient capabilities to leverage the given skill-based examples, it may not work better than Consistency@5 on these harder topics.”

3. “It is not clear that fixed examples are necessarily optimal for easy questions as performance is usually on par with randomly-selected examples”

On easy questions, we acknowledge that fixed examples usually only result in 1~2% advantages over random examples. However, on some domains (e.g., Intermediate Algebra and Counting & Probability, see Table 10 in Appendix D.2), randomly-selected examples can lead to 5% performance drop compared to fixed examples on the Qwen 1.5B model. In terms of average performance on the whole dataset, Figure 4 shows that using fixed examples on easy questions consistently outperformed random examples. Therefore, we believe that fixed examples are the optimal choice for easy questions in the scope of our study.

4. “the reference to cognitive theories as motivation for the learning strategy is bit surface-level and could be detailed further”

Thank you for the suggestion. The cognitive theories primarily serve as a motivational lens for studying adaptive in-context feedback in SLMs. We will include additional details and experiments in the next revision around our references to overthinking and adaptive teaching. We also refer the reviewer to our response 4 to Reviewer HGMM for related experiments on overthinking.

We would like to thank the reviewer for the careful review and their positive feedback on our study. Please find our responses to your questions below.

“Computational bottlenecks, due to requiring additional models like PRM and LLM.”

We refer the reviewer to responses 1 to reviewer Qc6H. In summary, we address overhead in two parts:

(1) the process reward model (PRM) is a general-purpose classifier trained once and applicable across math datasets, with alternatives like response length heuristics showing promising accuracy (up to 79.8%).

(2) AdaptMI relies on a one-time LLM labeling step to create the skill bank and is effective in itself. With additional overhead through repeated LLM use, AdaptMI+ can improve further.

“SkillBank’s construction seems to be impossible with SLMs only”

Our method uses a teacher-student framework where a relatively stronger model provides high-level skill annotations (not answers) based on metacognitive abilities. The teacher need not outperform the student; it only needs to be sufficiently capable of offering abstract skill-level guidance. In fact, in our response 2 to reviewer Qc6H, we show that our LLM scores 69% on MATH, below Qwen2.5-7B-Instruct. Thus, while we did not explore building SkillBanks with SLMs to keep our scope focused, we believe it an important direction to consider for future work.

“Ablation on using ICL examples rewritten to each SLM’s own reasoning style”

We conducted this ablation study on Qwen models of different scales. We:

(1) Prompted Qwen2.5-7B-Instruct to rewrite all the solutions in MATH by itself. We find that only the 7B model is capable of preserving key information while rewriting the solutions.

(2) Replaced all the in-context examples with the original questions and rewritten solutions.

(3) Tested the baselines and AdaptMI performance on all three models with the re-written ICL solutions.

As shown in table below, rewriting examples yields a ~1% accuracy gain for the 1.5B model, suggesting it struggles more with human-written examples. However, the gains are minimal for larger models.

“Statistics about length of skill-based examples and relation to (over)thinking they triggered”

First, we’d like to clarify that the fixed examples and skill-based examples are sampled from the same data distribution, i.e., the MATH or GSM8K training set. Hence, their average difficulty (and average length) should be similar to the output responses.

For easy questions in Level 1 and 2, the skill-based examples have similar lengths as fixed examples on average. Therefore, the models’ overthinking on easy questions is less likely to come from long skill-based examples. Rather, it may be caused by skill-based examples being too specific to one topic, as exemplified in our case studies in Section 4.1 (line 200-213).

More interestingly, we observe that for more difficult questions, their corresponding skill-based examples tend to be longer. This observation is likely because difficult questions often correspond to more advanced skills, such as solving equations or complex number calculation, and the associated ICL examples of these skills are longer.

“Distribution of difficulty level (as judged by the PRM) within each MATH subtask”

Below, we show a table with proportions of difficult questions for all the dataset domains. The proportion of difficult questions roughly mirrors (though slightly exceeding) the proportion of questions the model actually gets wrong (table 1 in main paper) in each domain.

We thank the reviewer for their detailed feedback on our work. Please find responses to your questions below.

“No details of the reward model”

Response: We would like to clarify that, in Section 3.1 (line 151-153), we specified the name, size, and training data of the reward model used, along with the prediction thresholds. Due to the space constraint, we put ablation studies on the reward model in Appendix B.1, where we different threshold values and compare process reward and outcome reward models. Please also check our responses to reviewer qoNx, where we add additional ablations on the effect of the thresholding functions. We will include them in our revision.

“Which skill-based examples resulted in the most improvement in accuracy?”

Response: As recommended by the reviewer, we conducted a detailed analysis to identify the skill categories where skill-based in-context examples are most beneficial. We grouped the questions based on their annotated skills and measured the average performance improvement of a small language model (SLM) when switching from fixed to skill-based in-context examples, separately for easy and difficult questions. For Qwen2.5-1.5B-Instruct model, we find that questions involving the skill "Perimeter and Area" exhibit the greatest performance gain for difficult instances. In contrast, for questions requiring the skill "Circle", performance notably declines for both easy and difficult variants.

“Can the "unnecessary information" or "cognitive overload" be quantified ?”

Response: In Figure 3, we showed that small models are cognitive overloaded by skill-based examples on easy questions, as their output lengths were much longer, but accuracy dropped. To understand “why” skill-based examples brought cognitive overload, we propose two explanations:

1. The skill-based examples are too long, which confuses the model. (This hypothesis was also raised by reviewer HGMM.)

2. The skill-based examples are too specific to one topic, which misleads the model or causes overthinking.

We already presented some case studies that indicate (2), in Section 4.1 (line 200-213). To test hypothesis (1), we calculated the average lengths of fixed examples and skill-based examples for questions in each difficulty level in Table B below.

We observe that for more difficult questions, their corresponding skill-based examples tend to be longer. Yet for easy questions in Level 1 and 2, the skill-based examples have similar lengths as fixed examples on average. Therefore, the models’ overthinking on easy questions is less likely to come from long skill-based examples. Rather, the cognitive overload may be caused by skill-based examples being too specific to one topic. We believe that this effect is harder to quantify and is best explained by our case studies in Section 4.1.

Table B: Average lengths of fixed examples and skill-based examples, in comparison with average output lengths of Qwen2.5-3B-Instruct, for questions in each difficulty level.

“Is there anything specific to pre-training of SLM that affects the performance of in-context skill based examples”

Response: The varied performance of different SLMs with AdaptMI and AdaptMI+ suggests that pre-training influences how models use skill-based in-context examples. However, without access to detailed pre-training information on the models, predicting the exact factors is difficult.

“Why are most gains in pre-algebra and the model struggles to make significant gains in geometry?”

Response: This likely reflects differences in pre-training; models may have seen more pre-algebra content than geometry. However, without detailed knowledge of the pre-training data, the exact cause is difficult to determine.