PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Self-improvement LLM Agentic System for ML Library Development

Genghan Zhang, Weixin Liang, Olivia Hsu, Kunle Olukotun

arXiv:2502.02534v2cs.CL

TL;DR

ML-library development for ASPLs is difficult and often has limited code examples, even though each DSA needs a library of operators. The paper introduces an adaptive self-improvement agentic system and evaluates it on a benchmark, reporting up to a 3.9× improvement over a baseline single LLM.

  • Problem

    ML libraries require expertise in both ML algorithms and ASPLs, while time-constrained development for evolving hardware leaves limited code examples and remains difficult even for experts.

    Source: §1 · pp. 1–2

  • Method

    An adaptive self-improvement algorithm integrates with an agentic system that filters correct answers, stratifies earned experiences by difficulty, and updates demonstrations.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    Up to 96% of benchmark tasks are solved, with up to a 3.9× improvement over a baseline single LLM.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The system supports experts developing ML libraries and produces resources that other systems can use.

    Source: §8 · p. 9

    When the task set exceeds the context length, better stratification and selection functions are needed to preserve experience quality.

    Source: §3.3 · p. 5

Abstract

from arXiv · show

ML libraries, often written in architecture-specific programming languages (ASPLs) that target domain-specific architectures, are key to efficient ML systems. However, writing these high-performance ML libraries is challenging because it requires expert knowledge of ML algorithms and the ASPL. Large language models (LLMs), on the other hand, have shown general coding capabilities. However, challenges remain when using LLMs for generating ML libraries using ASPLs because 1) this task is complicated even for experienced human programmers and 2) there are limited code examples because of the esoteric and evolving nature of ASPLs. Therefore, LLMs need complex reasoning with limited data in order to complete this task. To address these challenges, we introduce an adaptive self-improvement agentic system. In order to evaluate the effectiveness of our system, we construct a benchmark of a typical ML library and generate ASPL code with both open and closed-source LLMs on this benchmark. Our results show improvements of up to $3.9\times$ over a baseline single LLM.

1. Introduction

Developing ML libraries for evolving, architecture-specific languages requires expertise in both ML algorithms and the target language, often with few code examples. This paper proposes adaptive self-improvement for LLM agents and reports benchmark gains for generating STeP libraries.

1. Introduction

  • Each DSA needs an ML library of ASPL-written operators, and developing one requires expertise in both ML algorithms and the target ASPL.Library development composes low-level ASPL primitives into high-level ML operators.Source: §1 · p. 1
  • Evolving ASPLs and time-constrained library–chip co-design leave limited code examples, while library development is difficult even for expert programmers.The paper cites FlashAttention-3 appearing two years after H100 release and a 47% performance drop when adapting FlashAttention-2 from A100 to H100.Source: §1 · p. 1
  • Existing self-improvement approaches face limits in exploration or data efficiency, motivating agents that can reason from limited data.The cited reinforcement-learning approach requires hundreds of effective trajectories per problem.Source: §1 · p. 2
  • The proposed adaptive learning cycle filters high-quality answers, stratifies earned experiences, and adaptively selects demonstrations to improve agents.The system is designed to operate without human experts beyond ASPL designers who initially explain the primitives.Source: Figure 1 · p. 1, §1 · p. 2
  • Up to 96% of benchmark tasks are solved, with up to a 3.9× improvement over a baseline single LLM.Figure 2 compares Pass@n across single-LLM, agentic, and adaptive self-improvement agentic systems.Source: Figure 2 · p. 2, §1 · p. 2

2. Background

The background frames ML-library development as programming evolving DSAs through architecture-specific languages, then introduces STeP and its streaming dataflow model as the target setting. It highlights the portability and example scarcity challenges that motivate automation.

2.1. Architecture-specific programming languages

  • ASPLs provide low-level programming interfaces to DSAs through primitives and specialized functions, forming a bottom-up abstraction of the chip.The section also benchmarks eight operator groups within a common LLM layer.Source: §2.1 · p. 3, Figure 3 · p. 3

2.2. ML library development using ASPLs

  • Rapid ASPL evolution makes libraries difficult to port: functions using updated MMA instructions must be rewritten for each GPU generation.The paper proposes enhancing users’ ability to learn a given ASPL rather than automating only selected operators, since updates such as memory optimizations may affect any operator and new operators continue to appear.Source: §2.2 · p. 3

2.3. STeP for next-generation RDA

  • STeP is chosen as the target ASPL because it models streaming dataflow for next-generation RDAs and poses similar library-development challenges despite lacking executable examples.Its streams carry tokens between compute units, with control signals embedded in tokens; each stream can be consumed by at most one primitive.Source: §2.3 · pp. 3–4
  • STeP primitives divide into arithmetic operations on stream tokens and shape manipulations that change control tokens.Figure 4 illustrates a simplified MoE program and its execution on memory and compute units.Source: Figure 4 · p. 3, §2.3 · p. 4

3. Adaptive self-improvement learning

Adaptive self-improvement learns from its own successful task completions, filtering and stratifying earned experiences into demonstrations until tasks are solved or demonstrations run out. The process can extend to new tasks, but context length may require better selection strategies.

3. Adaptive self-improvement learning

  • The algorithm samples agents in parallel, filters successful answers, stratifies earned experiences, and updates demonstrations until tasks are solved or demonstrations are exhausted.It also removes completed tasks from the task set, adaptively allocating more test-time compute to harder tasks.Source: §3 · p. 4

3.1. Filtering high-quality answers

  • The filter groups correct answers by isomorphic syntax trees, randomly selects one representative per group, then keeps the shortest pure-code answer and records its success rate.The shortest-answer choice follows the Minimum Description Length principle, while random selection balances the risk of losing chain-of-thought comments.Source: §3.1 · p. 4

3.2. Stratification and selection

  • β bins experiences into m difficulty levels based on success rate and selects demonstrations incrementally; repeated steps are retained when newly solved tasks are easier than the demonstrations.Those lower-success-rate tasks have a better chance of yielding a correct answer when the number of samples is doubled; later iterations can also use fewer, higher-quality tokens.Source: §3.2 · pp. 4–5

3.3. Discussion

  • The process can incorporate new tasks, but if their number exceeds the context length, better stratification and selection are needed to preserve experience quality; a finite task set terminates.Each example takes about 500 tokens, so cutting-edge LLMs can handle hundreds more tasks.Source: §3.3 · p. 5

4. Agentic system organization

The system combines LLM agents, a code generator, and verifiers, using a structural intermediate representation to connect users and components. Its proposer and guardian agents generate implementations and address global affine-type constraints, while the surrounding tools generate and verify executable code.

4.1. LLM agents

  • The proposer generates candidate STeP implementations, while a guardian checks and corrects global affine-type violations that challenge causal, step-by-step LLM generation.The proposer prompt uses STeP references and usage patterns; the guardian consists of a fused gate and corrector.Source: §4.1 · p. 5

4.2. Code generator

  • The code generator turns implementations into self-contained pytestable Python scripts, letting agents focus on STeP semantics using familiar Python syntax.A functional simulator executes the embedded STeP specifications.Source: §4.2 · p. 6

4.3. Verifier

  • Verification checks functional correctness by comparing simulator outputs with PyTorch result tensors and checks affine types through static analysis of the STeP syntax tree.The simulator limits dynamism to what ML operators require and emulates stream execution with tensor computation and necessary control flow.Source: §4.3 · p. 6

4.4. Structual intermediate representation

  • The structural IR unifies interfaces between users and the system and between agents and the code generator, encoding programs compactly rather than merely reformatting prompts.For the compared task, users specify the operator and specialized functions without redundant glue text, saving tokens for more demonstrations.Source: §4.4 · p. 6

5. Benchmark

The benchmark covers common ML operators and specialized functions, and evaluates task completion with pass@k. Its tasks are designed to span operator groups and difficulty conditions, with a reference implementation for each task.

5. Benchmark

  • The benchmark contains 26 tasks across 8 groups of ML operators in common LLM model architectures.It is designed to cover diverse popular operators and specialized functions.Source: §5 · p. 6

5.1. Metric

  • Pass@k measures task completion from T tasks, n system samples, and c_i correct responses per task; pass@1 is expected success rate and pass@n is expected solvable-task share.The paper uses the two extremes, pass@1 and pass@n, to analyze performance.Source: §5.1 · p. 6

5.2. Benchmark Construction

  • Tasks are constructed from first principles with nearly balanced operator-group counts and an even spread of difficulty conditions.The construction includes tasks with or without reused variables and tasks requiring shape and arithmetic primitives; each task has a reference implementation.Source: §5.2 · p. 6

6. Experiments

Experiments find that adaptive self-improvement benefits from prioritizing hard-earned demonstrations, while mixed demonstrations can help discover useful hard examples. The system and its structural IR also show gains and trade-offs in base capability, semantic diversity, and affine-type correction.

6. Experiments

  • The evaluation uses the Section 4.1 agentic organization to assess base learning capability; prompts are formatted in YAML.OpenAI-o1 is benchmarked separately but excluded from later experiments to control test-time compute.Source: §6.1 · p. 7, §6 · p. 7

6.1. Analysis of adaptive self-improvement learning

  • Hard-only demonstrations form the Pareto optimum for all three models and deliver the largest learning-curve gains; fewer hard examples can beat larger mixed sets.Figure 7 tracks the adaptive sampling iterations and demonstration cycles.Source: §6.1 · p. 7, Figure 7 · p. 8
  • Mixed demonstrations can uncover later hard-only sets (DeepSeek-V3: H2 follows HM1/HME1) and save tokens even when they do not improve performance (GPT-4o, Llama).The examples show why mixed experiences remain part of adaptive sampling despite hard-only examples' stronger performance.Source: §6.1 · p. 7

Adaptive granularity affects token efficiency and peak

6.2. Ablation study of agentic system organization

  • Without self-improvement, the agentic system composes attention operators at over 50% Pass@1, including online softmax and memory-free streaming attention.The passage notes these specialized functions are challenging for ordinary programmers.Source: §6.2 · p. 7
  • Guardian raises Avg-Reuse Pass@n from 6/12 to 11/12 and Arith-Reuse from 5/8 to 8/8, but lowers Arith-Once Pass@1 from 0.685 to 0.663.The system's Pass@1 for Avg-Once remains unchanged.Source: Table 2 · p. 7, §6.2 · pp. 7–8
  • Structural IR's performance gains accompany increased sample diversity; Pass@n positively correlates with failure and overall diversity, while IR can hurt success diversity.The table analyzes correlations between semantic diversity measures and performance.Source: §6.2 · p. 7, Table 1 · p. 7

7. Related Work

Related work situates the paper among LLM self-improvement and specialized-task agentic systems. It distinguishes its task-level, AST-based scoring and adaptive sample selection, and highlights structural IR as a task-aligned interface.

7. Related Work

  • The paper positions its work at the intersection of LLM self-improvement and agentic systems for specialized tasks.The related-work discussion focuses on these two areas.Source: §7 · p. 8

7.1. Self-improvement learning for LLMs

  • Unlike token-level post-training, this method treats implementations as task-level actions and earned experiences as states, using AST analysis to score answers and adaptively prioritize high-scoring samples.Prior incorporation strategies include retrieval, reflection, and reward feedback.Source: §7.1 · pp. 8–9

7.2. Agentic system organization for specialized tasks

  • For specialized ML-library development, structural IR provides task-aligned interfaces between agents and tools within the agentic system.The paper situates this design among task-specific agentic organizations.Source: §7.2 · p. 9

8. Conclusion

The paper co-designs the learning process and agentic system to support complex reasoning with limited data in ML library development using ASPLs. Its methods implement ML operators while producing a self-improving agent.

8. Conclusion

  • The system co-designs its learning process and agentic organization around complex reasoning with limited data.The paper identifies ASPL-based ML library development as poorly automated.Source: §8 · p. 9
  • Its methods both implement non-trivial ML operators and produce a self-improving agent.Source: §8 · p. 9
  • The system supports experts developing ML libraries and offers resources for other systems.The paper envisions applying similar techniques to a broader class of tasks.Source: §8 · p. 9

Impact Statement

The paper describes potential productivity and access benefits from automating ML library development, alongside possible effects on programming work. It calls for ongoing discussion as the technology develops.

Impact Statement

  • The technique could enhance ML library developers’ productivity and accelerate development of more efficient ML systems.Source: Impact Statement · p. 9
  • The technique could democratize access to ML development tools and reduce technical barriers to entry.Source: Impact Statement · p. 9
  • Increased software-development automation could affect the nature of programming work and required skills, warranting ongoing discussion.The authors call for careful consideration as the technology develops.Source: Impact Statement · p. 9

A.1. STeP introduction

The STeP example presents a simplified mixture-of-experts program using typed streams and primitives for routing and accumulation. The accompanying text explains the operators and token types used in the example.

A.1. STeP introduction

  • Algorithm 2 is a STeP program example for a simplified mixture-of-experts module, with typed inputs and output.Its inputs include Buffer and Multihot values; the output is a Buffer.Source: §A.1 · p. 14
  • The type signature follows Haskell style, and three functions are supplied externally by the hardware.Buffer, Multihot, and Scalar are token types; Buffer and Multihot also carry shapes.Source: Figure 10 · p. 14, §A.1 · p. 14
  • Partition routes tokens to experts using an index stream, and Merge accumulates expert tokens using a shared index stream.Merge’s function updates an initialized state; Partition and Merge are used in pairs.Source: §A.1 · p. 14
  • Copy duplicates streams, Zip forms tuples, and Map applies a function to each input value.Promote, Repeat, and RepeatRef are shape-manipulation primitives.Source: §A.1 · p. 14

A.2. Benchmark details

This section details the benchmark’s task families and evaluation settings, reports additional system comparisons and code-quality results, and outlines possible extensions to other tasks and languages. It also provides implementation examples, including a Rowmax reduction over the innermost dimension and the Accum primitive’s specification and Python frontend definition.

A.2. Benchmark details

  • The benchmark spans task families including attention, three matrix-multiplication dataflow orders, positional-pairing variants, and LayerNorm and RMSNorm.Group 4 uses softmax(S)·V with S = QKT; the last three Table 3 tasks—index, expert, and etoe—come from MoE.Source: §A.2 · p. 15

A.3. Experiment settings

  • Experiments sampled 64 times per temperature in Section 6.1 and 64 times at temperature 0.7 on Claude 3.5 Sonnet in Section 6.2.Section 6.1 tested temperatures 0.4, 0.7, and 1.0; the settings also specify four model APIs, a 1024-token output limit, and GPT-4o seed 42.Source: §A.3 · p. 16

A.4. Additional experiments

  • The system completed each task in under 10 minutes on average, while a domain expert was unable to write one program within 48 hours during a pilot study.The system could perform explorations in parallel, whereas the expert had to use trial and error and accumulate experience sequentially; maintainability scores averaged MIwoc 102, MI 149, and MIcw 47.Source: §A.4 · pp. 16–18
  • OpenAI-o1 achieved similar performance to a single Claude-3-5-Sonnet at the cost of more tokens, while a single Claude-3-5-Sonnet proposer completed more tasks using fewer tokens.Table 4 compares the self-improvement system across models; the reported comparison here is between OpenAI-o1 and a single Claude-3-5-Sonnet proposer.Source: §A.4 · p. 16, Table 4 · p. 17
  • With token counts matched and a shared verifier, the agentic systems still performed better; GPT-4o supervised finetuning improved performance, but less than self-improvement.Self-improved models use the verifier throughout the process, while single and agent variants use it only as a final judge; the SFT study used 133 correct programs in its starting dataset.Source: Table 5 · p. 17, §A.4 · pp. 17–18, Table 6 · p. 18

A.5. Generalization

  • Adaptive self-improvement raised Claude-3-5-Sonnet Pass@n on AIME-2024 from 0.50 to 0.67.The dataset contains 30 challenging AIME 2024 problems; the proposed system organization requires defining task and output formats, building a verifier, and selecting an LLM or domain-tailored agents.Source: §A.5 · pp. 18–19

A.6. Additional explanation

  • Unlike Spiral, Ansor (Zheng et al., 2020), and AKG, the system uses LLMs to generate and refine ASPL ML-library code and adapt over time.The related-work discussion contrasts this approach with systems that rely on predefined optimization strategies and cost models.Source: §A.6 · p. 19

A.7. Code and prompt examples

  • The code examples illustrate structural-IR attention implementations and the use of Copy primitives to ensure a STeP stream is consumed only once.The appendix also shows Bufferize, Repeat, and Streamify shape-manipulation primitives and a guardian-agent prompt.Source: Figure 17 · p. 22, Figure 18 · p. 23, §A.7 · p. 25
  • Rowmax maps [B,N,E] to [B,N] by reducing over the innermost dimension.The example uses Accum with f32 inputs and outputs and a=3, b=1.Source: §A.7 · p. 26
  • Accum applies the max function on each dequeue and emits the accumulated value when it sees .Sx(x>=b), with b=1 in this example.The example requires seeing the whole vector to obtain the reduced value.Source: §A.7 · p. 26
  • Figures 21 and 22 show the Accum primitive’s specification in the STeP document and its definition in the Python frontend, respectively.Source: Figure 22 · p. 26, Figure 21

Related papers

Loading 2502.02534v2…