Procedural PRM — LoRA adapter for Llama 3.1 8B Instruct
A process reward model (PRM) fine-tuned to score whether a proposed next action is correct given a procedural workflow, the history so far, and a list of available actions. Answers Yes/No and at inference we read the logits to get a continuous score.
The goal is to be combined and wrapped around another LLM, even with a critic in the loop to check for factuality or harmfulness. The goal is to act like a specialized planning agent, which based on what it did in the past, can pick the correct future actions.
Training
- Base model: Llama 3.1 8B Instruct (4-bit quantized)
- Method: QLoRA via Unsloth
- Data: ~8000 step-level SFT records from procedural workflow extractions, not perfect data, but realistic. NOTE: there is also a version on my profile trained on 43% more data, but with duplicate traces.
- The training data includes near-miss cases where an action might sound plausible, but it's actually wrong. We include also the reason why it's wrong in the SFT records so the model can learn this
- Hardware: Single A100 on Google Colab, 3 epochs
Usage
It scores candidate actions — it does not generate text, except yes and no tokens for valid or not. it s not a classifier!!
Combine it with a frozen weights LLM (as capable as possible), tool use, and you will have a good planner in the loop.
Context
Built as part of a master's thesis on extracting procedural workflows from text and training small planning agents with the data generated by the extractor pipeline.
In short, get procedural documents, use the extraction agentic pipeline from my work, create SFT records deterministically and then re-train the PRM, acting as a planner or even can be seen as a critic which approves or not choices of a big LLM.
When enough data is generated, re-train the model.
