ChemAgent: Self-updating Memories in Large Language Models Improves Chemical Reasoning

We greatly appreciate your detailed feedback and insightful questions. Below, we address each concern, elaborating on improvements made and additional experiments conducted.

Q1. Why focus solely on chemistry applications?

Our initial focus on chemistry stems from its unique combination of challenges:

1. Complexity: Chemistry tasks require both rather long reasoning chains and the integration of domain-specific knowledge, such as formulas and constants.

2. Reasoning vs. Calculation: Unlike advanced mathematics, where highly complex and precise calculations (e.g., PDEs) are common, chemistry calculations are often simpler, making them ideal for validating our framework's reasoning capabilities.

3. Data availability: Chemistry benchmarks like SciBench offer rich, domain-specific evaluation datasets.

Q2. Is there anything chemistry-specific about the system?

The answer is nuanced, while ChemAgent's framework can smoothly extend to other domains, its memory library is tailored to chemistry during the library construction phase. For example:

• Execution Memory (M_e) stores solutions to chemical problems.
• Plan Memory (M_p) is built based on chemistry-specific task decompositions.
• Transitioning to a different domain (e.g., physics or mathematics) would require restarting the memory initialization process. Without this, applying chemistry-specific memories to other domains could degrade performance.

Q3. Performance in other domains

To address this, we tested GPT-4omini on other SciBench subdomains, including physics (CLASS) and mathematics (DIFF and STAT).

Results are shown below:

Key insights:

1. ChemAgent consistently outperforms CoT, demonstrating its flexibility in enhancing stepwise reasoning across domains.
2. However, direct reasoning performed best, echoing observations in the Table 3 of the original SciBench paper, which noted degraded performance when using tools like Python for some LLMs (GPT-4-Turbo for example).
3. This may indicates that distilled models (like GPT-4omini) may suffer from reduced reasoning and coding capabilities compared to GPT-4.

Q4. Clarify sources and methodologies in Table 1

Thank you for pointing out this oversight. In our revised submission:

• Few-shot + Python results were sourced from SciBench's benchmark paper (Table 3).
• Few-shot Direct Reasoning results were obtained from the StructChem paper.

We have updated the table's descriptions to accurately reflect these sources and methodologies, ensuring transparency.

Q5. Rationale for chosen benchmarks

We selected SciBench for three key reasons:

1. Development Set Advantage: SciBench provides a detailed dev set with clear solutions, streamlining our library construction phase.
2. Comparative Baseline: StructChem, a state-of-the-art method, is benchmarked on SciBench, providing a direct comparison.
3. Recency: SciBench was one of the most comprehensive and up-to-date benchmarks available when we began this research.

Q6. Performance on ChemBench

We evaluated ChemAgent on ChemBench's physical chemistry subset using GPT-4o.

Results are as follows:

ChemAgent demonstrated significant improvements, further validating its effectiveness. We have included a detailed example trajectory from ChemBench in Appendix H.

Q7. Predefined templates vs. ChemAgent's approach

Thank you for offering the valuable suggestions about our mistake on describing previous methods. We have modified the paper according to your advices.

About "predined templates": While ChemAgent uses predefined prompts to standardize task formats, it differs from template-based methods in key ways:

1. Agentic Workflow: ChemAgent dynamically adjusts the number and structure of subtasks based on task complexity, unlike rigid pre-designed templates.
2. Iterative Refinement: Through evaluation and refinement modules, ChemAgent can backtrack and revise strategies autonomously.

Q8. Missing citations and writing improvements

We have addressed the issues you raised:

1. Citations: Relevant works (e.g., [1-5], [7-11]) have been incorporated into the revised manuscript.
2. Writing: Ambiguous phrases (e.g., "predefined templates") and unclear sections (e.g., Table 1 descriptions) have been clarified. Changes are highlighted in red in the updated paper.

Q9. Ablation studies on the model used in the evaluation and refinement module

In Appendix B, Table 4, the last row demonstrates that using GPT-4 to evaluate ChemAgent's solutions and refining with GPT-3.5 outperforms the approach of using GPT-3.5 for both evaluation and refinement. However, in GPT-3.5 experiments, it still underperforms compared to ChemAgent operating without the evaluation and refinement module.

This experiment highlights that the effectiveness of the evaluation and refinement module is highly dependent on the self-correction capabilities of the refinement model.

Due to budget and computational resources constraints, we did not conduct further ablation studies to evaluate the impact of model selection on the performance of the evaluation and refinement module.

Thank you for kindly pointing out these issues.

Generalizability:

We sincerely appreciate your insightful feedback. In response, we have expanded our discussion on the generalizability of our method in the Conclusion and Appendix E.Limitation of the revised paper.

We would like to clarify that when changing domains, the core structure of our framework remains unchanged—only the content of the memory library is modified. However, during the experiments presented in Table 2 (Section 4.1), we observed that for problems with shorter reasoning chains, removing the plan memory sometimes improved performance. Thus, when applying our framework to less complex domains, such as middle school mathematics, a potential modification to our architecture could involve removing the plan memory.

Regarding Q3 (Experiment on Math and Physics)

Yes, we retrained the memory modules for each experiment. In Q3's case, we did Library Construction for 3 times.

Compare with the other methods evaluated in the chembench paper?

The ChemBench paper evaluates only two methods beyond direct prompting with an LLM: (1) answering based on the log probabilities of the answer tokens, and (2) ReAct. The results for physical chemistry, as reported in their official repository, are as follows:

These results are directly sourced from the ChemBench GitHub repository. (Note: The performance of zero-shot direct reasoning using GPT-4o in the first row has been updated from 33.33% to 50.00% compared to our previous response, as we initially tested it ourselves and obtained 33.33%. The value in the ChemBench repository, however, is 50.00%.)

As shown, our method outperforms all other approaches evaluated in the ChemBench paper.

We hope this response clarify your questions and provide additional insights. Thank you again for helping us improve our work. Feel free to reach out for any further clarifications or discussions.

We sincerely appreciate your insightful feedback and thoughtful questions about ChemAgent's memory management system. Please allow us to address your concerns comprehensively.

Q1. How to handle potential memory overflow? How do we dynamically optimize memory usage?

A1.
Generally, we have not focused on controlling the size of the memory library for the following reasons:

1. Our work primarily emphasizes the effectiveness of self-updating memories in reasoning tasks.
2. The dataset we use (SciBench) is relatively small, making memory overflow unlikely.

However, in Section 3.3 (page 8, paragraph 2: “To ensure accuracy and prevent target leakage, …, for any j < k.”), we discuss an example of managing memory usage. Specifically, we exclude certain memories during specific iterations by examining their similarity. For instance, in Section 3.3, memories with a similarity score higher than 0.88 with $\mathcal{P}_i$ are discarded. While this method controls memory usage, it does not directly optimize it.

A straightforward solution to handle memory overflow would be to discard memories with high similarity to existing ones. However, this approach may negatively impact the performance of ChemAgent. Addressing this trade-off remains an open question and is worth further exploration. We have added this limitation in the Limitations section of our paper.

---

Q2. What are the sizes of each memory module, and do they vary by type? What are the considerations behind these differences?

A2.
Yes, the sizes of memory modules vary by type. The size is influenced by whether ChemAgent is allowed to self-evolve during runtime (Section 3.3).

1. Without self-evolution during runtime:
   The sizes of plan memory and execution memory primarily depend on the size of the development set for a given dataset. For example, using the ATKINS dataset:

   • Execution memory: 8 MB
   • Plan memory: 3 MB
     The dev-set size is 16. As described in Section 2.4 and Figure 2, incorrect solutions are discarded during the library construction phase. Thus, datasets with the same dev-set size may have slightly different memory module sizes due to differences in discarded solutions.

   The memory statistics for various datasets are summarized below:

   Dataset	Dev-set size	Plan Memory	Execution Memory
   ATKINS	16	2.9 MB	8.1 MB
   CHEMMC	9	1.5 MB	5.3 MB
   MATTER	10	6.3 MB	4.1 MB
   QUAN	8	1.7 MB	3.7 MB

2. With self-evolution during runtime:
   In this case, the sizes of memory modules are approximately proportional to the number of tasks ChemAgent has encountered.

For both settings, the size of knowledge memory is dynamic because we do not store it. Instead, ChemAgent dynamically generates knowledge memory at runtime based on the specific problem it is solving.

There are no manually designed considerations behind the differences in memory module sizes. Execution memory is typically larger than plan memory because a single plan may include multiple sub-tasks, each contributing to execution memory.

---

Q3. How long does it typically take to build a knowledge library that achieves reliable performance?

A3.
The time required to build the library depends on the size of the dev-set and the inference speed of the backbone LLM. For example, with the ATKINS dataset and GPT-4:

• Approximately 6~7 hours are needed, mainly due to the latency of the GPT-4 API.

Using local models can significantly reduce this time.

The minimum dev-set size required to achieve reliable performance remains an open question and is outside the primary scope of our research.

---

Q4. Is there a mechanism to ensure the quality of data accumulated over time?

A4.
Yes, there are mechanisms in place to ensure data quality:

1. During the Library Construction phase: Memories that the LLM evaluates as incorrect trajectories are discarded.
2. During runtime experiments (e.g., Section 3.3): Only memories derived from correctly solved tasks are inserted into the library.

---

We hope these explanations clarify our approach and address your concerns. Please feel free to reach out for further discussion or clarifications.

We thank the reviewer for this insightful question. We want to clarify several key differences between ChemAgent and existing approaches:

Unlike existing prompting methods like Chain-of-Thought (CoT), Tree of Thoughts (ToT), or Buffer of Thoughts (BoT) that focus on improving reasoning through better prompting structures, ChemAgent introduces a novel dynamic memory system specifically designed for chemical reasoning. While methods like ToT maintain a tree structure for exploration, our approach maintains a structured knowledge repository that evolves with each problem-solving experience.

Compared to other memory-augmented approaches, ChemAgent incorporates domain-specific structure through its three distinct memory types (Planning Memory, Execution Memory, and Knowledge Memory). This specialized architecture allows for:

• Hierarchical decomposition of chemical problems
• Storage of verified solution strategies
• Dynamic updating of chemical domain knowledge

We have added a new comparison table that highlights these differences:

---

Could the authors elaborate on any unique aspects of ChemAgent's architecture or framework that differentiate it from similar systems in the literature?

The key architectural innovations of ChemAgent include:

• Domain-specific memory structure: Unlike general-purpose systems, ChemAgent's memory components are specifically designed for chemical reasoning:

  • Planning Memory captures high-level chemical problem-solving strategies.
  • Execution Memory stores detailed solution steps with chemical formulas and calculations.
  • Knowledge Memory maintains fundamental chemistry principles.

• Runtime memory updates: We have implemented a novel verification mechanism that ensures correct and useful information is added to the memory pool. Our experiments (Section 3.3) show this improves accuracy by 7% compared to systems without updating.

---

While the use of the SciBench dataset is commendable, could the authors clarify if the results obtained are robust across a wider range of chemical tasks? Have the authors considered evaluating ChemAgent on real-world chemical challenging reasoning tasks, such as optimizing experimental parameters or suggesting experimental protocols?

Thank you for this suggestion. We acknowledge this limitation. However, we note that comprehensive evaluation on experimental optimization tasks is currently limited by the lack of standardized benchmarks. This presents an opportunity for future work in creating such datasets.

The importance of chemical reasoning tasks is highlighted by recent works in drug discovery and materials science (citations added). Our framework provides a foundation for tackling these complex tasks through its ability to learn and refine chemical problem-solving strategies.

---

The authors should explore performance on tasks with varying levels of difficulty, such as those requiring multi-step reasoning or integration of data from multiple sources.

We agree with this suggestion and have added new experiments on other benchmarks (ChemBench's Physical Chemistry sub-dataset).

These results demonstrate our system's robustness across varying difficulty levels and different datasets while highlighting areas for future improvement.

Thank you for your insightful review. Below, we address each of your questions in detail.

---

Q1. Clarify the library construction phase.

A1.
Thank you for raising this question. We have added an example of the library construction phase in Appendix H. In summary:

• "Conditions" refer to the given parameters in a problem, such as "temperature is 27°C" or "volume is 10L."
• The "refine" step is essentially a mechanism designed to ensure that the extracted conditions are neither missing nor redundant.

The specific prompt used for this step is detailed in Appendix G.1, under the section titled REFLECT_PROMPT.

---

Q2. Can our method be applied to other scientific domains with similar reasoning demands?

A2.
Absolutely! Both you and reviewer h1D2 have highlighted this valuable point. Our method is indeed applicable to other scientific domains with similar reasoning requirements. Beyond minor adjustments to system instruction prompts and restarting the library construction phase, no significant modifications are necessary.

We have included additional experiments and analyses in our response to reviewer h1D2's Q2 and Q3, demonstrating broader applicability and presenting interesting insights.

---

Q3. Elaborate on the observed limitations in the evaluation and refinement module.

A3.
ChemAgent’s performance on the first test data (where 2θ is the correct answer) in ATKINS provides a representative example:
ChemAgent is able to correctly report 2θ as the answer when the evaluation and refinement module is not activated, however, when the evaluation and refinement module is enabled, ChemAgent reports θ instead, which means an unnecessary refinement process is triggered.

Additionally, as shown in Figure 7 (Error 3), there are some types of errors hard to correct, even when the errors are detected. In cases like Error 3, the refinement module often selects other incorrect or irrelevant memories again.

We attribute these limitations to:

1. The lack of relevant and useful memory.
2. The model's inherent difficulty in correcting its own reasoning errors without external guidance.

---

Q4. How does ChemAgent compare with models that integrate external tools, such as ChemCrow?

A4.
Integrating specialized tools is an excellent suggestion. However, most tools integrated into ChemCrow (e.g., Query_to_SMILES) are not well-suited for tasks in SciBench.

Our implementation, like ChemCrow, includes web search functionality. However:

• Since the questions in SciBench are sourced from textbooks and generally searchable online, we did not evaluate this feature.
• Additionally, web search introduces significant uncertainty into problem-solving.

Nonetheless, we appreciate this suggestion as a promising direction for future research.

---

Q5. How do we ensure the correctness of the generated code?

A5.
This is an excellent question. During initial experiments, we encountered issues with incorrect code generation. Our solution involves:

1. Retry Mechanism: If the generated code fails to execute, it is retried up to four times.
2. Fallback Strategy: If all retries fail, the model bypasses code generation and directly outputs its reasoning results.

In practice:

• Due to the simplicity of chemical codes, persistent code failures are rare.
• For mathematical problems, particularly those involving equation solving, repeated failures were more common.

We chose Python code over specialized computational tools to enhance the generalizability of our approach.

---

Q6. Why did we choose SciBench? And experiments on other datasets/benchmarks?

A6.
We selected SciBench primarily for its well-defined step-by-step solutions in the dev set, which facilitated the initialization of our memory library without requiring additional manual effort.

To evaluate broader applicability, we extended our library to test on ChemBench's Physical Chemistry dataset, achieving significant performance improvements:

These results underscore ChemAgent's effectiveness and generalizability across different datasets.

---

We hope these responses clarify your questions and provide additional insights. Thank you again for your thoughtful feedback and for helping us improve our work. Please feel free to reach out for further clarifications or discussions.