tarsur385/swev-prm-embeddings-qwen36-27b

SWE-bench Verified — PRM Precomputed Embeddings (Qwen3.6-27B)

Precomputed Qwen/Qwen3.6-27B pooled embeddings of (state, action) pairs from SOTA mini-SWE-agent trajectories on SWE-bench Verified. These are the cached embeddings used to train a retrieval / process reward model (PRM) for best-of-N action selection (pick the candidate action per step that maximizes cos(proj_state(state), proj_action(action))).

Embeddings are frozen 5120-d vectors; the PRM trains two lightweight projection heads on top of them with InfoNCE + hard negatives (large batches give many in-batch negatives — this is why the cached-embedding setup works where end-to-end LoRA re-embedding does not).

Contents

File	Shape / type	Description
train/state_embeddings.pt	[401509, 5120] fp32	pooled embedding of the agent state (message history) at each step
train/action_embeddings.pt	[401509, 5120] fp32	pooled embedding of the action string at each step
train/metadata.json	dict (samples list)	per-sample metadata (trajectory_id, step_idx, task_id/instance_id, domain, reward, ...), aligned row-for-row with the embeddings
val/*	[94210, 5120]	same, held-out split
slim_train.json / slim_val.json	list	slim per-pair records (task_id, trajectory_id, action, model, step_idx, reward) — enough to regenerate hard negatives

Train ≈ 401,509 pairs · val ≈ 94,210 pairs · dim 5120.

Split

Task-disjoint train/val split: no SWE-bench instance_id appears in both splits (same split as the companion SWE-bench Verified TRM). Prevents trajectory/instance leakage.

How it was built

• State: mini-SWE-agent message history → chat template (add_generation_prompt=False) → Qwen3.6-27B vLLM pooling embed, head-truncated to 32K tokens.
• Action: the <function=bash>{"command": ...}</function> command string → raw vLLM embed.
• Hard negatives (not stored — fully reproducible): for each positive, real actions emitted by failed (resolved=False) trajectories on the same instance_id (primary) or same repo (backoff). Since an action embedding is a pure function of its text, each negative string maps directly to a row in action_embeddings.pt — no re-embedding needed at any negative count.

Usage sketch

import torch, torch.nn.functional as F
S = torch.load("train/state_embeddings.pt")   # [N,5120]
A = torch.load("train/action_embeddings.pt")  # [N,5120]
# train proj_state, proj_action (Linear-GELU-Linear -> 512) with InfoNCE over the batch + hard negs
# score(state, action) = cos(F.normalize(proj_state(s)), F.normalize(proj_action(a)))

Companion dataset: TRM trajectory embeddings tarsur385/swev-trm-trajectories-25models.