COGTRL: Training LLMs for Scientific Discovery Assistance using Cognitive Traces via Reinforcement Learning

The half of the paper nobody writes down

A research paper records the method that worked. It almost never records the thinking that produced it: the constraint that ruled out the obvious route, the alternative that failed, the reason step three had to come after step two. That cognitive process is exactly what a scientist working toward a goal under constraints wants from an assistant, and it is exactly what is missing from the literature LLMs are trained on. COGTRL's premise is that a model should be trained to think its way to a method, not just to write one: given a research goal and a list of constraints, it emits a short cognitive trace before every methodological step, and reinforcement learning rewards the two together.

Two rewards and an uplift

A trajectory is an alternating sequence (t₁, s₁, t₂, s₂, …, tₙ, sₙ): each cognitive trace tᵢ explains why the next step should advance the goal under the constraints, and each step sᵢ is conditioned on everything before it. A large reasoning model, OpenAI o3-mini, acts as the reward model and scores every trace and every step from 1 to 5 on six dimensions each:

• Trace quality — goal and constraint integration, scientific and mechanistic reasoning, causal logic and actionability, information density, scientific accuracy and consistency, and explicit uncertainty and trade-offs.
• Step quality — alignment with the research objectives and constraints, scientific plausibility, innovation and novelty, testability, feasibility and scalability, and impact potential.

The piece that makes it work is the uplift reward, R_uplift = R_step · σ(R_trace − α) with α = 0.6. A trace is not rewarded on its own merits; it is rewarded in proportion to the step it precedes, gated by whether the trace itself clears a quality threshold. The policy therefore cannot collect points for eloquent explanations that don't change the method. A structural reward enforces the strict <Trace_i> … <Step_i> alternation, and the total R_total = R_step + γ·R_uplift + λ·R_struct is optimized with GRPO: sample a group of trajectories per input, normalize each reward against the group to get an advantage, and update with the clipped objective plus a KL term toward the SFT reference.

Interleaved beats think-first

Two generation formats were trained: think-first, where all cognitive traces come before the full methodology, the way current thinking models behave; and interleaved, where each step follows its own trace. Interleaving wins consistently, by 4.30 points on AI and 4.99 on Materials Science averaged over both models. Scores are the six step-quality dimensions summed and scaled to 100, judged by two independent agentic evaluators and averaged over three runs:

That puts the trained 3B models in 70B territory. Zero-shot with the same judges, Llama-3.3-70B scores 52.27 / 54.00 and Qwen-2.5-72B scores 50.47 / 54.57 on AI / Materials Science; frontier models still lead (GPT-5.4 at 60.27 / 64.60, Claude Opus 4.6 at 63.87 / 62.53). Against vanilla GRPO on the same base model, COGTRL adds 4.12 points on average; against every non-COGTRL baseline, 7.85. The gain is spread across all six quality dimensions, largest on testability (+4.55), impact potential (+4.55) and scientific plausibility (+4.50). Swapping GRPO for Dr. GRPO gives the same picture, so the effect isn't tied to one RL algorithm.

Traces have to earn their keep

Three analyses show the traces are doing causal work rather than decorating the output:

• Remove the uplift reward and quality drops by 5.38% on average (Llama 52.63 / 55.64 → 50.65 / 53.20; Qwen 51.13 / 53.17 → 48.35 / 48.96). Rewarding traces and steps separately is not enough; the coupling matters.
• Swap the traces. Teacher-forcing fixed high-quality steps, their per-token probability is highest under the aligned trace, lower with no trace, and lowest with a mismatched trace from another example in the same domain (Llama on AI: 0.632 vs 0.527 vs 0.473). The benefit comes from trace–step alignment, not from extra context or formatting.
• Ask for traces at inference. Prompting a COGTRL model to emit traces lifts its methodology quality over the same model answering without them (Llama: 49.50 / 51.12 → 52.63 / 55.64).

No collapse on general reasoning

Specializing a small model usually costs it something elsewhere. Supervised fine-tuning does exactly that here: averaged over AIME, AMC23, GPQA-Diamond, HumanEval, MMLU-Pro and OlympiadBench, SFT with traces halves Llama-3.2-3B's score (21.07 → 10.38) and cuts Qwen-2.5-3B's from 32.01 to 22.48. COGTRL holds the line and edges up: 21.37 for Llama and 33.31 for Qwen, with gains on AMC23 for both models. Training the scientific reasoning did not come at the expense of the general kind.

What the experts saw

Four annotators, two chemistry PhD students, one CS PhD student and one CS master's student, blind-ranked 300 methodologies (50 problems per domain, three systems each) with a Kendall's W of 0.70. COGTRL's output was preferred in 71.42% of pairwise comparisons, and the reasons were consistent: completeness, organization and causal structure. In materials science that meant concrete synthesis and characterization workflows and appropriate use of molecular dynamics and DFT; in AI it meant grounding a method in conceptual reasoning before spelling out training and evaluation. The paper's case study is a photothermal self-healing elastomer under three coupled constraints: the zero-shot model proposes a generic workflow, the GRPO model introduces sophisticated components without linking them, and the COGTRL model decomposes the problem into constraint-driven stages whose synthesis, characterization and modeling steps each cite the constraint they satisfy.

Building the data

No dataset of scientific cognitive traces existed, so the paper built one. An agent pipeline extracts the goal, constraints and step-by-step methodology from 4,990 arXiv papers across six domains (Physics 1,311 · CS 1,225 · Math 1,116 · AI 583 · EE 570 · Biology 185); GPT-4o then drafts three candidate traces per step at different temperatures and a critic ensemble of GPT-4o, Gemini-2.5-Pro and o1 keeps the best. Those teacher traces seed the SFT initialization only; RL explores its own. Evaluation uses two expert-extracted test sets, 50 NeurIPS 2024 papers (AI, in-domain) and 50 papers from the MatDesign benchmark (materials science, out-of-domain), all from 2024 to sit past the models' knowledge cutoff, with substantial inter-annotator agreement (weighted κ = 0.67). Judging is rubric-based rather than reference-matching and is done by two independent agentic evaluators, Codex with GPT-5.4 and Claude Code with Claude Sonnet 4.6, that are distinct from the o3-mini reward model, so the policy can't simply learn its judge; their verdicts match human expert preferences 77.14% of the time.

Why it matters

Small open models can become genuinely useful discovery assistants if you train the thinking rather than the transcript. The mechanism is general: wherever a model produces intermediate rationales, reward them by what they cause downstream rather than by how they read, and the rationales stop being decoration.

Resources

Paper: arXiv:2608.30109 · EMNLP 2026 · Code: github.com/shri071/CoGTRL