Adaptive Self-improvement LLM Agentic System for ML Library Development

We thank Reviewer 4wDa for the positive comments and helpful feedback. We were encouraged that the reviewer enjoyed reading the paper, found our ideas novel, and took the time to review the appendix. We will include all the discussions and results below in the revised version.

Run-time performance measurement

”Q1: Have you measured the actual run-time performance of any of these automatically generated kernels on real or simulated hardware?”

We manually translated the generated implementation of Task 2 in Figure 16 of our paper to a simulator built on top of DAM-RS, which models the streaming behavior of each STeP primitive and assumes every operation and specialized function takes one cycle. This task implements softmax(S)@V where S and V are of shape [n,m] and [m], respectively; since S and V are streamed sequentially, the cycle count should scale with n*m. Simulation results match this expectation. Detailed result: https://anonymous.4open.science/r/ICML2025-rebuttal-4D6B/fig.png

Generality of the framework

”Q2: Do you anticipate any unique challenges in applying this approach to more mainstream, well-established languages like CUDA/HIP or to CPU vector intrinsics?”

Our approach has two parts: adaptive self-improvement learning and agentic system organization. The learning process is broadly applicable; the challenges lie in tailoring agentic systems to other languages.

Mainstream languages like CUDA, HIP, and CPU vector intrinsics exhibit global properties such as arbitrary memory access, data layout sensitivity, and side effects. Similarly, STeP enforces a global affine type constraint. Our framework addressed this using a guardian agent that detects and corrects affine type violations. This concept generalizes: domain-specific guardian agents can monitor and enforce global properties of various languages, adapting the STeP solution more broadly.

A second challenge is that LLMs may lean toward surface-level patterns in mainstream languages due to their existence in training data, potentially missing more optimal or novel transformations. As shown in Section 6.2, our structural IR can increase sample diversity and thus boost the LLM agentic system performance. Extending this, structural IRs and tailored code generators can guide LLMs toward more creative solutions beyond conventional patterns.

Code maintainability and complexity

”Q3: Have you evaluated the readability or maintainability of the final code solutions?”

”Metrics such as lines of code, the number of shape transformations, or specialized instructions would help quantify the difficulty.”

Maintainability statistics. We assessed code maintainability using two metrics: maintenance index without comments (MIwoc) and with comments (MI) [1]. The comment weight (MIcw) is defined as MI - MIwoc and falls in [0, 49); MI > 85 indicates good maintainability. Using all correct programs from our best model (self-improved agentic Claude Sonnet), we recorded the top MIwoc and MI per task. The mean MIwoc is 102, MI is 149, and MIcw is 47—indicating well-commented, maintainable code.

Complexity statistics. We used the same set of programs as the maintainability statistics to measure complexity. We measured all three metrics the reviewer suggested:

• Lines of code: Counted via primitive calls (excluding comments/blank lines)
• Shape transformations: Counted by use of Promote, Repeat, RepeatRef, and Flatten primitives
• Specialized instructions: Counted as the number of specialized functions in task descriptions

Across all completed tasks:

Detailed result: https://anonymous.4open.science/r/ICML2025-rebuttal-4D6B/tab.md

Scalability

”…it is unclear how the system would scale on libraries that are much bigger or that require 50–100 times more operators.”

As discussed in the “Larger scale evaluation potential” section of our response to Reviewer qrAB, current LLMs still have the capacity to self-improve over hundreds more tasks. If the number exceeds the context length, better stratification and selection functions are needed to preserve experience quality within the context window limit.

Related work

”... TVM and Spiral, and more recent papers on end-to-end auto-tuning code generators, may be relevant. Additionally, a broader set of reflection-based LLM coding pipelines, such as Reflexion, and Tree-of-Thought could be cited or compared.”

We thank the reviewer for providing more relevant work. We will incorporate these papers into related work and give a more thorough discussion of how our work improves the results of them.

Reference

[1] https://www.verifysoft.com/en_maintainability.html

We thank Reviewer 6GnU for the constructive comments and helpful feedback. We are encouraged that the reviewer found our improvement non-trivial. Below, we respond to the raised concerns.

Challenge of this programming task

”It is not clear why this programming task should be so challenging.”

We appreciate the question. While Claude Sonnet achieves 70% on our benchmark, this does not imply that ML library development using ASPLs is easy. As a real-world example, as we pointed out in the paper, it took the community two years post-H100 release to optimize attention—a key LLM operator—to ~70% peak performance. As an analogy, although Claude Sonnet scores 81.7% Pass@1 on HumanEval-Mul (Table 6 in [1]), generating code from language instructions remains a hard problem. Our benchmark represents only a subset of the broader challenge: implementing transformer operators using a Python-embedded, side-effect-free ASPL. Many other ML operators and ASPLs involve more complex semantics. We plan to expand the benchmark with more difficult tasks to push system capabilities further.

Generality of the framework

”Also the generality of the framework is unclear - is it only for that particular language?”

We thank the reviewer for raising this important point. While our evaluation focuses on one language, STeP, we argue that the framework is general. As discussed in Section 2.1, we identify two essential features of ASPLs—primitives and specialized functions—and show in Section 2.3 how STeP embodies them. Due to space constraints, we refer the reviewer to the "Generality of the framework" section of our response to Reviewer 4wDa for the challenges and solutions of applying our framework to other languages.

Comparison fairness

”… This seems unfair as their agentic systems perform a lot of extra computation and have access to tools (like the verifier).”

As the reviewer pointed out, differences in tool access and computational load might have influenced the outcomes. Therefore, we conducted an experiment that aligned both aspects.

For computation fairness, we matched the token count of the single model with the agent and self-improved models by resampling. All model variants (single, agent, self-improved) have access to the same verifier so it is fair; the difference lies in how the verifier is leveraged. Self-improved models incorporate it throughout the process, while others use it only at the end as a final judge. We chose Claude Sonnet and GPT-4o as base models. Below is the result:

Therefore, our agentic systems still perform better under this fair setting. Since we also agree aligned comparisons can provide a more comprehensive view, we will add these results in the revised version.

Finetuning

”…, it is unclear, if a system fine-tuned say on samples that got filtered by the verifier or in some other way, would not outperform the proposed system”

We conducted supervised finetuning (SFT) using GPT-4o and found it improved performance, but less than our self-improvement approach.

Since we do not know the exact SFT algorithm of OpenAI service for FLOPs matching, we tried our best to favor the SFT method. We began with the same 133 correct samples from all completed tasks used in the first iteration of self-improvement. Different from self-improvement which only picks 1 correct program per completed task, we picked all 133 programs to form the training dataset of SFT. We created three SFT datasets with varying prompt compositions:

• 133 (base prompt+question+answer)
• 133 (question+answer)
• 17 (question+answers deduplicated via AST)

Each dataset was used to train a separate SFT model. After that, we sampled each model on all the uncompleted tasks. Below is the result:

We appreciate the reviewer’s suggestion and will include these results in the revised version.

Paper organization

”The paper should more clearly carve out early on what the contribution to the ML field are”

”The claims like "we do human style learning with some ref" are too brief and vague… “

We thank the reviewer for these helpful suggestions. In the revised version, we will emphasize our contributions to ML more clearly in the introduction and better define human-style learning in the discussion.

Reference

[1] Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C. and Dai, D., 2024. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437.

We thank Reviewer hNb4 for the positive comments and helpful feedback. We are encouraged to hear the reviewer found our experiments and demonstrations clear and convincing. We also appreciate the reviewer’s careful review of the entire appendix content for prompt details.

Other agentic methods for ML library development using ASPLs

”…, are there other existing Agentic systems/workflows that can be applied to the current task?”

We thank the reviewer for raising this important point. As discussed in the paper, the tight timeline of the library-chip co-design process highlights the need for better automation. The community has begun to explore agentic solutions to this challenge in parallel [1, 2, 3]. Our adaptive self-improvement learning can enhance these efforts by making better use of correct samples. Additionally, existing systems often struggle with evolving language features, whereas our method, designed for new languages, adapts naturally to such changes.

Other fields and scenarios potential

”Can some experiments be designed in the future to prove the reliability of the system in other fields and scenarios?”

”I hope the authors can discuss the potential of the proposed system in other fields and scenarios.”

The proposed system can be extended to other scenarios that require complex reasoning with limited example data and well-defined evaluation metrics. We outline the general recipe below.

As shown in Figure 1 of our paper, the agentic system organization is constructed in three main steps. First, system designers define the format of both the task and its expected output. Once the format is specified, the next step is to build a verifier for the task. With the format and verifier in place, designers can either use a single LLM or design LLM agents tailored to the domain—similar to how we handle the type constraints of the STeP language. After completing these three steps, the task can be handed over to our system, which will automatically carry out adaptive self-improvement learning.

The adaptive self-improvement learning system also exposes several tunable hyperparameters which will be helpful when the results are not satisfactory. The most direct control is the number of parallel sampling. Users can also adjust the adaptive granularity parameter m for experience stratification. Additionally, domain-specific filtering heuristics—such as the minimal code length heuristic used by us—can be incorporated to further guide the learning process.

We also conducted an experiment on the AIME-2024 dataset which contains 30 challenging problems from the American Invitational Mathematics Examination (AIME) 2024. We applied our adaptive self-improvement learning to the Claude Sonnet base model and increased Pass@n from 0.50 to 0.67. This demonstrates the potential capabilities of our system on other tasks.

Since we agree with the reviewer that demonstrating potentials in other domains can benefit the community, we will add the recipe and results in the revised version.

References

[1] https://developer.nvidia.com/blog/automating-gpu-kernel-generation-with-deepseek-r1-and-inference-time-scaling/

[2] Lange, R.T., Prasad, A., Sun, Q., Faldor, M., Tang, Y. and Ha, D., 2025. The AI CUDA Engineer: Agentic CUDA Kernel Discovery, Optimization and Composition.

[3] Ouyang, A., Guo, S., Arora, S., Zhang, A.L., Hu, W., Ré, C. and Mirhoseini, A., 2025. KernelBench: Can LLMs Write Efficient GPU Kernels?. arXiv preprint arXiv:2502.10517.

We thank Reviewer qrAB for positive comments and helpful feedback on our work. We are encouraged to hear the reviewer found the task and method to be innovative and the experimental results to be promising.

Larger scale evaluation potential

”The benchmark dataset, consisting of only 26 tasks across 8 groups, is relatively small and could benefit from a larger scale evaluation.”

We thank the reviewer for highlighting the potential benefits of large-scale evaluation. We briefly addressed this point in Section 3.3 of the paper, and we will expand the discussion in the revised version. New tasks typically involve new ML operators and new hardware specialized functions. These can be incorporated into the existing task pool and handled using the same adaptive self-improvement learning process by selectively sampling only the new tasks. In our current experiments, the longest prompt is approximately 14k tokens, with each example averaging around 0.5k tokens. Given Claude Sonnet’s 200k-token context window, there is capacity to include hundreds of additional tasks.

Related work

”The paper could strengthen its literature review by incorporating more recent work on agent self-improvement, such as ADAS[1], AFLOW[2] to better position its contributions.”

We thank the reviewer for pointing out two relevant works—ADAS [1] and AFLOW [2]—that can enhance our literature review. The revised version will cite both papers.

Human programmer comparisons

”The inclusion of human programmer comparisons would provide valuable context and better demonstrate the practical significance of the system's achievements.”

We appreciate the reviewer’s suggestion. Our system completed each task in under 10 minutes on average. In contrast, during our pilot study, a domain expert was unable to write a single program within 48 hours, as they had to do trial-and-error and accumulate experience sequentially. Our system, by comparison, can perform these explorations in parallel. We agree with the reviewer that the comparison to human programmers offers valuable insight, and we will include these results in the revised version.

In the future, we can also collaborate with HCI researchers to conduct more extensive experiments on the time and effort required by human programmers versus our system, aiming to better understand usability and cognitive load.

References

[1] Hu S, Lu C, Clune J. Automated design of agentic systems[J]. arXiv preprint arXiv:2408.08435, 2024.

[2] Zhang J, Xiang J, Yu Z, et al. Aflow: Automating agentic workflow generation[J]. arXiv preprint arXiv:2410.10762, 2024.