Kishan25/neuron-kd-qwen2.5-coder-0.5b-critique-mlx-4bit
Neuron KD — Critique Model (Qwen2.5-Coder-0.5B, MLX 4-bit)
A 0.5B on-device model that explains why Python code fails and suggests a fix, given the code and the result of actually running its tests.
This is the code-review sibling of `neuron-kd-qwen2.5-coder-0.5b-mlx-4bit`, which writes solutions. Both power the Code Lab in the Neuron iOS app.
The design decision that matters
A 0.5B model is bad at judging whether code is correct — that's static analysis, and it can't reliably simulate execution in its head. Trained to be a judge, it collapses to answering "Looks correct" to everything:
So this model isn't asked to judge. The host app runs the tests with a real Python interpreter and passes the result in the prompt. The model's job is only to explain a known failure and repair it — and it is trained on prompts that contain that execution feedback, so the format matches at inference time.
Training on prompts without the feedback and adding it later does not work: the earlier model offered a fix just 5 times out of 40 even when explicitly told the code had FAILED. Matching train and inference prompts took that to 37 out of 40.
Measured results
Held-out MBPP test split, never trained on. Repairs are verified by executing the model's suggested fix against the real asserts — no LLM judge, no human scoring.
Across 5 sampled runs over 976 candidates (temperature 0.2):
Best-of-3 works because the host executes each candidate fix and keeps the first that passes. Retries use a higher temperature (0.6); attempt 3 still contributed 12 of 80 successes, so the retries find genuinely different approaches rather than re-rolling the same idea. Average cost: 2.22 generations per repair, since it stops early on success.
The honest limitation: roughly half of broken solutions get no verified fix. The intended UX is to show an explanation only in that case, never to present an unverified fix as an answer.
Prompt format
Execution feedback must be included, exactly where shown:
Problem: {problem description}
Your solution must pass this test:
{asserts}
Here is a candidate solution:
{candidate code}
The solution was executed against the tests.
Result: FAILED
Failing test: {the assert that failed}
{error output}
Review the candidate solution. Respond in exactly this format:
Verdict: <"Looks correct" or a short phrase describing the issue>
Explanation: <1-2 sentences>
Suggested fix:
<a corrected Python function, or "None needed" if the solution is correct>When the tests pass, replace the feedback block with All tests passed.
Training
- Base: Qwen2.5-Coder-0.5B-Instruct
- Teacher: Qwen2.5-Coder-32B-Instruct, given the execution result and asked to explain it — so teacher critiques never contradict reality (0 verdict/execution mismatches in 282 sampled)
- Data: 1,608 candidates across 420 MBPP problems — reference solutions, teacher solutions, AST mutants (swapped comparisons, off-by-one constants, flipped operators), and empty stubs. Every candidate labelled by executing it, never by opinion. 762 passing / 846 failing.
- Objective: hybrid KD — cross-entropy on teacher text plus KL against the teacher's top-5 token distributions (α=0.5, T=2.0)
- Method: LoRA r=32, α=64, on attention projections; 3 epochs
- Trained alongside solution-writing examples so that skill isn't overwritten — a larger-capacity run (r=64 + MLP targets, 6 epochs) overfit: training loss fell to 0.056 while held-out pass@1 dropped below the untouched base model.
Usage
pip install mlx-lm
mlx_lm.generate --model Kishan25/neuron-kd-qwen2.5-coder-0.5b-critique-mlx-4bit \
--prompt "$(cat prompt.txt)" --max-tokens 300 --temp 0.2// MLX Swift
let container = try await LLMModelFactory.shared.loadContainer(
configuration: ModelConfiguration(
id: "Kishan25/neuron-kd-qwen2.5-coder-0.5b-critique-mlx-4bit")
)Scope
Trained on MBPP-style single-function Python exercises with assert-based tests. It is not a general-purpose code reviewer, and it should not be the thing that decides whether code is correct — run the tests.
