reasoning-memory

Procedural Knowledge at Scale Improves Reasoning

This repository contains the minimal, end-to-end pipeline for reproducing the paper results: generate a procedural knowledge datastore, build retrieval indices, run retrieval, perform model rollouts with retrieved subroutines, and filter the samples to output the final metrics.

Overview

The reproduction pipeline has five steps:

• Step 1: generate subquestions and subroutines from raw corpora.
• Step 2: embed the datastore and build retrieval indices.
• Step 3: retrieve procedural knowledge for evaluation datasets.
• Step 4: run model rollouts using retrieved knowledge.
• Step 5: merge rollout samples and report metrics.

Setup

Data

1. Raw corpora go in raw_corpus/. Follow the download guides in:
   • raw_corpus/openthoughts/README.md
   • raw_corpus/nemotron/README.md
2. Evaluation datasets go in data/. Follow:
   • data/README.md

Environment

Install the repo requirements from the repo root. You might need to install the appropriate PyTorch version before this step.

pip install -r requirements.txt

If you plan to use vLLM for generation or query decomposition, start OpenAI- compatible servers and set the URLs in the step scripts below.

Running Experiments

All steps are run from the repo root. Each .sh script is a thin wrapper with defaults you should edit for your hardware, models, and datasets.

Step 1: Generate the Datastore

bash step1_corpus_generation.sh

Key settings (edit in step1_corpus_generation.sh):

• SOURCE: openthoughts or nemotron.
• BACKEND: vllm or matrix (https://github.com/facebookresearch/matrix).
• TEACHER_MODEL_ALIAS and VLLM_BASE_URL: model alias and OpenAI-style URL.
• --max-rows: set to -1 for full processing.

Outputs are written to reasoning_memory_datastore/{source}/ as examples.jsonl and datastore.jsonl.

Step 2: Build the Retrieval Index

bash step2_index_generation.sh

Key settings (edit in step2_index_generation.sh):

• SOURCE: matches the datastore source from Step 1.
• MODEL_NAME: embedding model (default ReasonIR).
• SHARDS: set to an integer to write datastore.emb-*.partXX.npz.

Outputs are written to reasoning_memory_datastore/{source}/.

Step 3: Run Retrieval

bash step3_run_retrieval.sh

Key settings (edit in step3_run_retrieval.sh):

• TASK, SPLIT: evaluation setup.
• RETRIEVER, MODEL_NAME, TOPK, Q_TYPE: retrieval configuration.
• AUX_MODEL_ALIAS, AUX_BASE_URL: required when Q_TYPE=self_decomp (the default setting usedi n the paper where the model generates the query itself).

Outputs are written to retrieval_logs/{source}/. When indices are sharded, the script merges into a final retrieval.json.

Step 4: Run Rollouts

bash step4_run_rollout.sh

Key settings (edit in step4_run_rollout.sh):

• MODEL_PATH, PORTS: model and OpenAI-style ports for vLLM servers.
• HINT_FILE: retrieval output from Step 3.
• TOPK_HINTS, REPEATS_PER_HINT: rollout sampling controls.

Outputs are written to eval_results/.

Step 5: Merge Rollout Samples

bash step5_merge_rollout_samples.sh

Key settings (edit in step5_merge_rollout_samples.sh):

• RESULT_FILE: rollout JSONL from Step 4.
• TOKENIZER_MODEL_NAME: tokenizer used for length-based filtering.
• TOP_K, POOL_SIZE, SEED: aggregation controls.

This step reports metrics to stdout.

Note: we use the number of tokens for length-based filtering, which is slightly different from the paper setting (number of words) and thus the threshold might need to be configured differently.

Data Layout

• raw_corpus/: raw OpenThoughts or Nemotron corpora.
• reasoning_memory_datastore/: datastore outputs + indices.
• data/: evaluation datasets.
• retrieval_logs/: retrieval outputs from Step 3.
• eval_results/: rollout outputs from Step 4 and inputs to Step 5.

License

reasoning-memory is FAIR Noncommercial Research License licensed, as found in the LICENSE file.