PaperlayerView exact revision on arXiv

Source-linked AI summary

MUSE: Multi-Level Memory with Selective Runtime Injection for Scientific Coding Agents

Haoran Jin, Defu Lian, Xiting Wang, Hong Xie

arXiv:2610.38332v1cs.LGcs.AI

TL;DR

Scientific coding memory can miss implementation lessons whose relevance emerges only during execution. MUSE combines global, approach-level, and implementation-level memory with selective runtime injection and review advice. It scores highest on all six main-comparison metrics, though evaluation uses one solving and memory-construction backbone.

  • Problem

    Scientific coding memory may fail to distinguish shared checks, problem-specific approaches, and implementation pitfalls that become relevant only during execution.

    Source: §1 · pp. 1–2

  • Method

    MUSE builds global, approach-level, and implementation-level memory from historical attempts, supplying the first two at initialization and selectively injecting implementation experience with review advice during coding.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    MUSE achieves the highest scores on all six metrics in the main comparison across four benchmarks spanning numerical solving, kernel optimization, and scientific software repair.

    Source: §1 · p. 2, §5 · p. 8

  • Takeaways & Limitations

    Memory transferred from KBX to KGB increases correctness and acceleration rates relative to the memory-free agent.

    Source: §5 · p. 8

    Experiments use ZCode with GLM-5.3-flash and frozen memory banks, leaving benefits across other agent backbones and interfaces and continual memory updates unexamined.

    Source: Limitations · pp. 8–9

Abstract

from arXiv · show

Historical experience can guide scientific coding at multiple levels, but task-level retrieval may miss implementation lessons that become relevant during execution. We introduce MUSE, a memory management framework combining multi-level memory with selective runtime injection for scientific coding agents. MUSE builds global, approach-level, and implementation-level memory from historical trajectories. At initialization, the agent receives global memory and problem-matched approach-level memory. During execution, current interactions are used to retrieve implementation-level memory. An LLM gate selects relevant experience and generates specific review advice, or skips injection. Across four benchmarks spanning numerical problem solving, kernel optimization, and scientific software repair, MUSE outperforms retrieval and memory baselines on all evaluation metrics. Ablations support the multi-level design and gated injection, while transfer experiments show that memory remains useful across kernel benchmarks. These results show that multi-level memory and selective runtime injection enable effective reuse of experience in scientific coding. Code is available at https://github.com/hr-jin/MUSE.

1 Introduction

Scientific coding agents must coordinate software operations with scientific requirements, and memory can preserve reusable reasoning and implementation knowledge from prior attempts. MUSE addresses mismatched memory levels and implementation pitfalls that emerge during execution, and leads the evaluated methods across four benchmarks.

1 Introduction

  • Scientific coding memory must distinguish shared checks, problem-specific approaches, and local implementation experience, while targeting pitfalls that task descriptions may not reveal.Relevant queries and memory keys should match code operations, implementation choices, and errors rather than broad task context.Source: §1 · pp. 1–2
  • MUSE combines global, approach-level, and implementation-level memory with selective runtime retrieval and review advice for scientific coding agents.Global memory and problem-matched approaches are supplied at initialization; selected implementation experience can be injected during coding or withheld.Source: §1 · p. 2
  • MUSE scores highest on all six metrics across four benchmarks, while ablations favor adding memory levels and selective injection with review advice.Memory transferred from one kernel benchmark retains gains in correctness and acceleration on another.Source: §1 · p. 2

2 Related Work

Prior agent-memory work organizes histories, extracts reusable knowledge, and retrieves contextual guidance; scientific-memory work also preserves research artifacts and reusable implementations. MUSE combines these ideas for scientific coding tasks.

2 Related Work

  • Agent-memory approaches organize histories into tiers or linked notes, extract transferable principles, and retain execution procedures or lessons from failures.Examples include distillation, skill and workflow memories, ACE's playbook, and Agent KB's workflow guidance with execution-level lessons.Source: §2 · p. 2
  • Scientific-memory systems preserve research knowledge or artifacts, while other systems reuse verified implementations, kernel pairs, or optimization skills with applicability conditions.MUSE focuses on reusing experience from completed scientific coding tasks to guide an existing agent on new problems.Source: §2 · p. 3
  • Runtime-memory methods retrieve or adapt experience from current context, and some systems selectively decide when to inject reminders or guidance.MUSE combines selective runtime injection with shared global lessons, problem-matched approaches, and implementation experience retrieved for the current interaction.Source: §2 · p. 3

3 MUSE

MUSE reuses completed scientific coding attempts through three memory levels: shared global lessons, problem-matched approaches, and implementation experience retrieved during coding. An LLM gate selects experience that supports a specific check or revision and can withhold injection.

3 MUSE

  • MUSE uses global memory shared across tasks, approach-level memory matched to the problem, and implementation-level memory matched to the current interaction.The agent receives global and approach-level memory at initialization; an LLM gate assesses runtime experience against the task and visible execution history.Source: §3 · p. 3

3.2 Memory Construction

  • Source attempts yield approach lessons and implementation lessons that record concrete choices, their rationale, outcomes, and conditions of applicability.Both successful and failed attempts are retained, so reusable approaches and potential pitfalls can inform later work.Source: §3.1 · p. 4, §3.2 · p. 4
  • Global memory is consolidated and reviewed against source implementations, outcomes, and applicability conditions before the complete banks and indexes are frozen for testing.Consolidation can add, replace, or retire entries while preserving their conditions and rationale; review resolves conflicts and corrects explanations.Source: §3.2 · p. 4

3.3 Memory Initialization

  • At initialization, MUSE supplies the complete global bank and retrieves approach lessons and source material by problem similarity.Sources already represented in the initial references are excluded from runtime indexes.Source: §3.3 · p. 4

3.4 Selective Runtime Injection

  • Runtime retrieval uses the current interaction and separate indexes for implementation lessons, historical submissions, and errors to form candidate records.Merged candidates exclude previously injected sources, and at most four reach the gate.Source: §3.4 · pp. 4–5
  • The gate checks whether source conditions hold and remain unverified in the visible history, then can specify a revision or check suited to the ongoing solution.A selected experience and its review note enter context together; an empty selection leaves context unchanged, while a rejected source remains eligible later.Source: §3.4 · p. 5

4 Experiments

Across four scientific coding benchmarks, MUSE achieves the highest scores on all six metrics, with ablations and transfer experiments examining how its memory levels and runtime gate contribute. Case studies illustrate how global, approach-level, and implementation-level memories guide kernel work.

4.1 Experimental Setup

  • About 30% of each benchmark’s problems supplied shared source trajectories for memory construction; the banks were frozen before evaluation on the remaining problems.Evaluation covered PDE solvers, two kernel benchmarks, and scientific repository changes.Source: §4.1 · p. 5

4.2 Main Results

  • 92.2% pass rate on PDEAgent-Bench versus 87.8% for EvoScientist; MUSE ranks highest on all six Table 3 metrics and raises acceleration rates on both kernel benchmarks.The table reports results across four scientific coding benchmarks; its Accel. measure counts correct tasks with at least a 1.01× speedup over the benchmark reference.Source: Table 3 · p. 6, §4.2 · pp. 6–7

4.3 Ablation Study

  • On 451 PDE test tasks, pass rate rose from 44.8% for Vanilla to 82.7% with L0, 89.4% with L0+L1, and 92.2% with gated L2.L0 alone remained below the strongest memory baseline, while L0+L1 exceeded it.Source: §4.3 · p. 7, Table 4 · p. 7
  • Gate selection with review advice scored higher than direct similarity-based injection on every benchmark; KBX correctness rose from 66.4% without the gate to 79.5% with it.The no-gate control retained all three memory levels and retrieval rules while injecting the highest-similarity eligible record directly.Source: §4.3 · p. 7, Table 6 · p. 7

4.4 Cross-Benchmark Memory Transfer

  • On KGB, transferred KBX memory achieved 62.1% correctness and 25.0% acceleration, versus 44.3% and 8.6% without memory; its scores were close to target-extracted memory.Memory extracted from KGB achieved 66.4% correctness and 24.3% acceleration.Source: §4.4 · p. 7, Table 5 · p. 7

4.5 Case Study

  • In KGB logaddexp2, global L0 guidance to upcast half-precision inputs led MUSE to convert both inputs to FP32, pass evaluation, and achieve a 1.048× speedup.The Vanilla implementation passed FP16 operands to tl.exp2 and failed to compile.Source: §4.5 · p. 7, Figure 2 · p. 8
  • For KGB FP8 conversion, L1 retrieved an encoding approach ranked first among 63 extraction tasks; MUSE passed all 32 tests with a 1.380× speedup, while Vanilla passed none.MUSE encoded 1.0 as 0x38, unlike Vanilla’s ordinary integer cast to 0x01.Source: §4.5 · p. 8, Figure 2 · p. 8
  • For KBX normalization, runtime retrieval selected a two-pass variance method ranked 13th by initial similarity and outside the five L1 references; the revised target passed with a 1.884× speedup.The agent changed variance computation to accumulate squared deviations from the mean; the historical source achieved 1.818× speedup.Source: §4.5 · p. 8, Figure 2 · p. 8

5 Conclusion

MUSE combines global, approach-level, and implementation-level memory with selective runtime injection and review advice. The conclusion reports strong benchmark results, favorable ablations, and gains when transferring memory between kernel benchmarks.

5 Conclusion

  • MUSE supplies global lessons and problem-matched approaches at initialization, then selectively injects implementation experience with review advice during execution.The framework combines three memory levels with an LLM gate.Source: §5 · p. 8
  • MUSE achieves the highest scores on all six metrics in the main comparison.Cumulative solver ablations favor all three memory levels, and gate ablations favor selection with review advice over direct similarity-based injection.Source: §5 · p. 8
  • Memory transferred from KBX to KGB increases correctness and acceleration rates relative to the memory-free agent.Source: §5 · p. 8

Limitations

The evaluation is limited in model and benchmark scope, and uses frozen memory banks rather than examining continual updates.

Limitations

  • Experiments use ZCode with GLM-5.3-flash for solving, memory construction, and gating; additional backbones and agent tool interfaces remain unevaluated.The authors note these evaluations would help characterize variation in benefits and computational costs across agents.Source: Limitations · p. 8
  • Evaluation covers four benchmarks in numerical solving, kernel optimization, and scientific software repair, with transfer tested from KernelBenchX to KernelGenBench-MS.Additional transfer directions and scientific coding domains, including tasks with different libraries and numerical requirements, remain future work.Source: Limitations · pp. 8–9
  • The evaluation uses frozen banks and does not examine continual memory updates.Continually accumulated trajectories could be used to study consolidation, conflict resolution, and retirement of outdated experience as software environments evolve.Source: Limitations · p. 9

A.1 Research Assets

The appendix reports research-asset documentation, licensing, and human-participant status.

A.1 Research Assets

  • Table 2 gives extraction and test set sizes, while Appendix B provides English prompts for memory construction and injection.Source: §A.1 · p. 11
  • PDEAgent-Bench is distributed under CC BY 4.0,1; the cited passage also names KernelBenchX and KernelGenBench but does not state their licenses.Source: §A.1 · p. 11
  • The study involves no human annotation, human evaluation, or research with human participants; evaluation uses public-benchmark verifiers on compute servers.SWE Science tooling and documentation use MIT, while task repositories, dependencies, and scientific materials retain their upstream licenses.Source: §A.1 · p. 11

A.2 Implementation and Evaluation

The implementation uses ZCode with GLM-5.3-flash for solving, memory construction, and gating, alongside embedding-based retrieval and benchmark-specific software environments.

A.2 Implementation and Evaluation

  • ZCode with GLM-5.3-flash handles solving, memory construction, and gating, with high reasoning effort and a 128,000-token per-response limit for task-agent and gate requests.Task wall-clock limits are two hours for PDEAgent-Bench, four hours for each kernel benchmark, and three hours for SWE Science.Source: §A.2 · p. 11
  • Retrieval uses Qwen3-Embedding-8B with 4,096-dimensional float32 unit vectors and cosine similarity; runtime queries are neither chunked nor truncated.Source: §A.2 · p. 11
  • Initialization uses Top5 retrieval, while each of the four runtime routes returns Top1 when its query is available.Source: §A.2 · p. 11
  • The PDE environment uses Python 3.11.15, DOLFINx 0.10.0, and PETSc 3.25.4; both kernel environments use Python 3.12.13 and CUDA 12.8 PyTorch builds.KBX uses PyTorch 2.9.1+cu128 and Triton 3.5.1, while KGB uses PyTorch 2.9.0+cu128 and Triton 3.5.0; SWE Science uses task- and verifier-specific containers.Source: §A.2 · p. 11

A.3 Computational Resources

The evaluation used a 128-core AMD EPYC server with eight A100 GPUs, and MUSE’s memory construction and retrieval required no model-weight training or fine-tuning.

A.3 Computational Resources

  • The server had two AMD EPYC 7742 processors with 128 physical CPU cores and eight 40 GB NVIDIA A100-SXM4 GPUs.Source: §A.3 · p. 11
  • Each kernel evaluation used one GPU, with reference and candidate measurements under the same GPU lock; GLM-5.3-flash inference used hosted APIs.Source: §A.3 · p. 11
  • Total API expenditure was approximately US$1,500 across preliminary experiments, memory construction, and evaluation, while memory construction and retrieval used no weight training or fine-tuning.Source: §A.3 · p. 11

B Prompts

The appendix presents prompts for extracting, consolidating, and reviewing scientific memory, then selecting and formatting runtime injections. The prompts specify source handling, gate behavior, and the information passed to the gate.

B Prompts

  • The prompt appendix defines “two-level” as L1 and L2, “training” as memory-extraction sources, and the stored decision field as an implementation choice.Source: B Prompts · p. 12

B.1 Per-attempt Memory Extraction

  • Per-attempt extraction builds memory from one frozen historical attempt and treats its trajectory, submission, evaluation, and feedback as evidence rather than a new task attempt.Source: Figure 3 · p. 12, §B.1 · p. 12

B.2 Global Memory Consolidation

  • Global memory is built incrementally from frozen memories and historical evidence, with original evidence consulted to verify claims.Source: Figure 4 · p. 13, §B.2 · p. 13

B.3 Global Memory Review

B.4 Runtime Gate System Prompt

  • The runtime gate chooses at most one candidate, checking visible work and applicability; advice may suggest an approach or a concrete recheck, since claims of correctness are not verification.Source: Figure 6 · p. 15, §B.4 · p. 15

B.5 Runtime Gate User Message

  • The gate input template includes the task summary, prior public trajectory, current interaction, and retrieved candidates.Source: Figure 7 · p. 16

B.6 Runtime Injection Message

  • The injection template presents specific review advice followed by the selected historical experience.Source: Figure 8 · p. 17

Related papers

Loading 2610.38332v1…