neuralnoise.com

[WIP] Benchmarking Local LLMs Against Coding Agent Harnesses

Tue, 28 Apr 2026 01:00:00 +0100

I’ve been running a small benchmark, harness-bench, that pairs local LLMs (served via llama.cpp’s llama-server) with agent harnesses (Aider, Claude Code, OpenCode, Pi, Qwen CLI) on 16 software-engineering tasks across Python, PyTorch, JAX, C, C++, Rust, and SQL. Each (model, harness, task) cell is sandboxed: the agent only sees a scratch workspace/ and grading is done by a hidden test.sh that the agent never sees. The current sweep is 17 model-quants × 5 harnesses × 16 tasks = 1360 runs on a single M3 Max / 128 GB laptop.

(The benchmark repository itself – task prompts, hidden graders, raw agent traces – is private, to keep the task set out of training corpora. If the prompts and graders ended up in a future model’s pretraining data, the benchmark would stop measuring what it claims to measure. Aggregated results, the per-cell CSV, and the plotting source for this post are public; the tasks are not.)

These are very preliminary findings – some of the questions below probably deserve a careful re-run before I’d trust the rankings to two decimal places. But the headline patterns are stable enough across the Q4 and Q8 sweeps that I think it’s worth writing them down.

The 16 tasks span seven languages and a wide spread of difficulties:

The hardest end (pt3_rope_gqa, jax1_complex_lp, pt7_prompt_blend, pt6_generate_cached, rs1_arena, pt5_logit_lens) is what discriminates the top tier of (model, harness) cells from the rest – everyone passes sql1_recursive and p2_shortest_path.

Q1 – What are the best (model, harness) combinations?

The single perfect cell on the matrix is Qwen3.6-27B (UD-Q4_K_XL) + pi – 16/16 at ~207 s/task. It is the only combination across all 85 tested cells (Q4 sweep + Q8 sweep) that clears the full benchmark.

If you care about throughput, gpt-oss-120b (MXFP4) + pi scores 15/16 at ~34 s/task – about 6× faster than the top combo for one extra failure, and the fastest member of the 15/16 tier overall. If you want a “dense-feeling” mid-size model, Qwen3.6-35B-A3B (UD-Q4_K_XL) + qwen hits 15/16 at ~108 s/task.

Eight cells cleared 15/16 or better:

Combo	Pass	Avg time
`Qwen3.6-27B` UD-Q4_K_XL + `pi`	16/16	207 s
`gpt-oss-120b` MXFP4 + `pi`	15/16	34 s
`Qwen3.6-35B-A3B` UD-Q4_K_XL + `qwen`	15/16	108 s
`Qwen3.6-35B-A3B` Q8_0 + `claude`	15/16	244 s
`gemma-4-26B-A4B-it` UD-Q4_K_XL + `opencode`	15/16	307 s
`Qwen3.6-27B` Q8_0 + `qwen`	15/16	395 s
`gemma-4-31b-it` Q4_K_M + `pi`	15/16	422 s
`Qwen3.6-27B` Q8_0 + `opencode`	15/16	462 s

pi shows up 3 times in the 15/16-or-better tier; qwen and opencode twice; claude once; aider not at all.

A clearer way to see the same information is to plot every cell in the matrix on (avg time, pass rate). The pareto frontier is what you actually care about as a user:

Qwen3.6-27B + pi (top-left blue circle, ~207 s) is the only 100% point. gpt-oss-120b + pi (purple diamond at the far left) gets you ~94% in ~34 s – about 6× faster, one extra failure. The bottom-right cluster of cells past 700 s is mostly aider and opencode paying for verbose tool-use scaffolding without recovering accuracy.

Q2 – Ranking between models

Marginalising over the 5 harnesses (Q4 sweep, 5 × 16 = 80 cells per model):

Model	Quant	Pass	Rate	Avg agent time
`Qwen3.6-27B`	UD-Q4_K_XL	66/80	82.5%	474 s
`gemma-4-31b-it`	Q4_K_M	65/80	81.2%	582 s
`Qwen3.6-35B-A3B`	UD-Q4_K_XL	64/80	80.0%	215 s
`gpt-oss-120b`	MXFP4	62/80	77.5%	67 s
`Qwen3-Coder-Next`	UD-Q4_K_XL	56/80	70.0%	162 s
`gemma-4-26B-A4B-it`	UD-Q4_K_XL	56/80	70.0%	397 s
`Qwen3.5-35B-A3B`	Q4_K_M	55/80	68.8%	123 s
`Qwen3.5-27B`	Q4_K_M	55/80	68.8%	402 s
`gpt-oss-20b`	Q4_K_M	45/80	56.2%	96 s
`Qwen3-Omni-30B-A3B-Instruct`	Q4_K_M	27/80	33.8%	113 s

A few things stand out. The dense Qwen3.6-27B and the (relatively) new gemma-4-31b-it are roughly tied at the top, but the MoE Qwen3.6-35B-A3B (3B active in a 35B body) is essentially as accurate at less than half the wall-clock cost. gpt-oss-120b at MXFP4 buys you another ~3× speedup for ~3% accuracy. And Qwen3-Omni-30B-A3B-Instruct – the omni-modal sibling of Qwen3.6-35B-A3B – is the worst model in the sweep at 33.8%, suggesting that omni-tuning has a real cost for code-only agentic use.

Plotted against total parameter count (point size = active parameters, colour = family):

Two things to take from this. First, size alone is not the right axis – the 120-billion-parameter gpt-oss-120b lands below the 27-billion-parameter Qwen3.6-27B, and the omni-modal 30B point is the lowest in the chart. Second, active parameters do most of the work: in the 25-35B band the dense models (Qwen3.6-27B, gemma-4-31b-it) and the 3B-active MoEs (Qwen3.6-35B-A3B) sit on essentially the same accuracy line, and the (much smaller) gpt-oss-120b point sits ~5 points below. Family is the strongest single signal – the omni-tuned point is dramatically off-trend.

Q3 – Ranking between harnesses

Marginalising over the 10 Q4 models (10 × 16 = 160 cells per harness):

Harness	Pass	Rate	Avg agent time
`pi`	123/160	76.9%	163 s
`qwen`	120/160	75.0%	191 s
`claude`	106/160	66.2%	306 s
`opencode`	102/160	63.8%	271 s
`aider`	100/160	62.5%	384 s

pi is the strongest harness on this benchmark – and also the fastest of the top three. The accuracy gap between pi/qwen and claude/opencode/aider is a clear ~10-15 percentage points, big enough that I wouldn’t want to call it noise. (Caveat: Claude Code in particular is being talked to via the Anthropic-compat shim against a local llama-server, which is not the configuration its prompt scaffolding was tuned for.)

The (model × harness) breakdown shows that the harness ranking isn’t uniform across models – some models are very harness-sensitive (gpt-oss-120b collapses with opencode, Qwen3.6-27B collapses with aider), others are relatively flat (Qwen3-Omni, gpt-oss-20b). Note the single 100-pass cell (top-left) and the gemma-4-26B-A4B-it + opencode outlier (94 pass rate against opencode’s overall mediocre numbers) – a clear hint that any cross-harness conclusion based on a single cell is shaky.

Q4 – Do models or harnesses cheat by reading the hidden tests?

This is the question I was most curious about going in, and the answer is more interesting than I expected.

I grepped every agent.log for references to test.sh, tasks/.../test.sh, or “grader” – both in the model’s own reasoning and in tool calls (Read, Bash, Glob). Per harness:

Harness	Mentions test.sh / grader	Actually read or ran the hidden grader
`pi`	0	0
`qwen`	5	0
`claude`	7	0
`aider`	10	0
`opencode`	23	14

For qwen, claude, aider, the references are all benign: the model paraphrases the prompt’s own mention of “the hidden grader” or “the grader checks …” while reasoning out loud. None of them tried to glob, read, or invoke the grader.

opencode is a different story. Across 14 distinct (model, task) cells, opencode either ran bash /Users/pminervi/workspace/harness-bench/tasks/<id>/test.sh <workspace> directly to get a pass/fail signal, or read the contents of the hidden test.sh file with its Read tool – the latter giving the model the literal reference implementation and tolerance constants. A representative trace from Qwen3.5-27B-Q8 + opencode on pt3_rope_gqa:

AssertionError: output does not match either RoPE layout
  (half=0.9700, interleaved=0.8935)
→ Read ../../../../../tasks/pt3_rope_gqa/test.sh
"Let me look at the test file to understand exactly what it's expecting"
← Edit rope_gqa.py

Of those 14 peeking cells, 13 passed. So this isn’t a marginal effect: when opencode peeks at the grader, it usually wins. Now, two important caveats:

The hidden test.sh files live outside the workspace (under tasks/<id>/test.sh). The benchmark is sandboxed only by convention – the agent has full filesystem access to ~/workspace/harness-bench/, so all five harnesses could technically peek. Only opencode actually does, by default. That’s a property of the harness, not of the underlying model.
The two distinct behaviours (running the grader vs. reading the grader source) carry very different weights. Running bash test.sh workspace to iterate on a pass/fail signal is closer to “generous test-time compute” than to leaking labels; reading the test.sh source – which contains reference numerical tolerances and sometimes inlined reference implementations – is straight-up data leakage.

I do not see evidence of pi, qwen, claude, or aider attempting either of these things in any of the 640 (= 4 harnesses × 160) Q4 cells I inspected. Nor do I see models hard-coding test inputs into their solutions – the pt6_generate_cached task has an explicit anti-cheat (it counts embedding-layer token visits to catch implementations that ignore the KV cache and just re-run the model), and pass rates there look behaviourally consistent.

So: yes, one harness cheats by default, and it’s opencode. The implication is that any cross-harness comparison that includes opencode is mildly contaminated – specifically, ~14% of opencode’s passes (14 of 102) coincide with grader peeking. If I subtract those, opencode’s pass rate drops from 63.8% to roughly 55%, which would put it last. I’m planning to re-run with a stricter sandbox (chroot or container) before drawing harder conclusions.

Q5 – Does 4-bit vs 8-bit quantisation make a difference?

For the seven sub-50B models I re-ran the full sweep at Q8_0. The headline is that Q8 is a slight net regression, not an improvement, on this benchmark:

Model	Q4 Pass	Q8 Pass	Δ	Q4 jax1	Q8 jax1
`Qwen3.6-35B-A3B`	64/80	66/80	+2	5/5	3/5
`Qwen3.6-27B`	66/80	65/80	-1	2/5	4/5
`gemma-4-31b-it`	65/80	59/80	-6	0/5	0/5
`Qwen3.5-27B`	55/80	52/80	-3	2/5	0/5
`Qwen3.5-35B-A3B`	55/80	51/80	-4	5/5	3/5
`Qwen3-Coder-Next`	56/80	55/80	-1	0/5	2/5
`gpt-oss-20b`	45/80	52/80	+7	3/5	2/5
7-model total	406/560	400/560	-6	17/35	14/35

A few observations:

The aggregate is essentially flat (-6 / 560 = -1%) – well within the noise I’d expect from re-running a stochastic agent loop. Q4 is not obviously degraded relative to Q8 at this scale.
Q8 is also strictly slower. Qwen3.5-27B averages 402 s/task at Q4 vs longer at Q8; the bandwidth cost of 8-bit weights on Apple Silicon is real.
Per-model, the swings are noisy and not all in the same direction. Only gpt-oss-20b clearly benefits from Q8 (+7, the biggest single-model swing in either direction). gemma-4-31b-it loses the most at Q8 (-6, mostly distributed across small per-task regressions, including a 13/16 → 8/16 collapse for gemma-4-31b-it + opencode).
Quantisation effects show up most clearly on the hardest, numerically-finicky tasks. The two cleanest examples are jax1_complex_lp (ComplEx + filtered MRR) – where Q8 gives 14/35 vs Q4’s 17/35, mixed at the model level – and pt7_prompt_blend, the only task where Q8 cleanly beats Q4 (51% vs 42%). The dominant pt7 failure mode is a structural one (returning unmixed branch logits instead of probability-space-averaged log-probs), and Q8 appears to help models stay coherent over the longer reasoning chain that pt7 needs.

The summary I’d give on quantisation is: Q4_K_M (or UD-Q4_K_XL) is the right default on Apple Silicon for this kind of agentic-coding workload. You give up almost nothing on the 16-task average and you get 1.5-2× the throughput. The two exceptions are gpt-oss-20b (small, gains real accuracy at Q8) and any workload that leans heavily on long-chain numerical correctness like pt7_prompt_blend.

What’s still WIP

This post is tagged WIP for a reason. Loose ends I want to close before I’d call any of this final:

Stricter sandboxing so opencode’s grader-peeking behaviour doesn’t contaminate the harness ranking. I want to know what opencode actually scores when it can’t see test.sh.
Multiple seeds per cell. Right now each (model, harness, task) is one run. Variance on the borderline cells (the difference between 13/16 and 14/16) is probably comparable to the harness-to-harness gap I’m citing.
Task coverage. Sixteen tasks is small, and pt3_rope_gqa and rs1_arena are doing a lot of the work in separating the top tier. I’d like to add a few more genuinely-hard tasks, especially in C/C++ and SQL where the catalog is currently thin.
Vendor-hosted baselines. The benchmark is currently entirely local. Running claude against the actual Anthropic API and qwen against Alibaba’s API would let me decompose how much of the harness gap is “the harness scaffolding” vs “the local model is weaker than what the harness was designed against”.

The harness-bench repo (task prompts, hidden graders, raw agent traces) is private for the contamination reasons noted at the top – happy to share access on request for collaborators or for replication, just drop me a line. The aggregated per-cell CSV and the seaborn source for every figure in this post are public and live alongside the post itself, under scripts/plots/harness_bench/ – run python3 scripts/plots/harness_bench/make_plots.py to regenerate.

Deep Research MCP

Tue, 21 Apr 2026 01:00:00 +0100

To make my life a bit easier, I built deep-research-mcp, a small Python agent that exposes several “deep research” backends through a single Model Context Protocol server (Anthropic, 2024). It allows Claude Code, Codex, Gemini CLI, or any MCP client to fire off long-running research tasks against whichever backend the user prefers.

What the server exposes

The MCP server exposes deep_research, research_with_context, and research_status. The first kicks off a task; the second resumes one after an optional clarification round; the third polls. Everything else – provider selection, timeouts, clarification models, system prompts – lives in ~/.deep_research.

Architecture

flowchart LR Client["Claude Code / Codex / Gemini CLI"] -->|MCP| Server["FastMCP server"] Server --> Agent["DeepResearchAgent"] Agent --> Clar["Clarification"] Agent --> Instr["Instruction builder"] Agent --> Backend{{"Backend"}} Backend --> OAI["OpenAI Responses / Chat Completions"] Backend --> Gem["Gemini Deep Research"] Backend --> DrT["DR-Tulu /chat"] Backend --> ODR["Open Deep Research (smolagents)"]

DeepResearchAgent optionally runs a clarification pass (triage the query, ask follow-up questions, enrich it), optionally rewrites the query into a longer research brief, and then delegates to one of four backends behind a common interface. All backends return the same normalised record – report, citations, reasoning steps, task id, execution time – so the MCP tools and the CLI do not care which one ran.

Backends

OpenAI Responses API (default). Uses a Deep Research model such as o4-mini-deep-research with built-in web_search_preview and code_interpreter tools, background mode, and polling. Follows the reference pattern in OpenAI’s cookbook (OpenAI, 2025a; OpenAI, 2025b).
OpenAI Chat Completions. Same backend file, different code path (api_style = "chat_completions"). No built-in tools, no polling – just a blocking chat call. This is the escape hatch for any OpenAI-compatible endpoint: Perplexity’s Sonar Deep Research (Perplexity, 2025), Groq, Together, Ollama, vLLM, or llama.cpp’s llama-server.
Gemini Deep Research. Implemented against the Interactions API, with Google Search and URL context as built-in tools (Google, 2024).
DR-Tulu. Allen AI’s open research agent, called over its /chat endpoint; the client-side integration is intentionally thin (AI2, 2025).
Open Deep Research. A self-contained smolagents stack with a text browser and search tools, using LiteLLM as the model layer (Roucher et al., 2025).

Why MCP?

Plugging a deep-research tool into Claude Code takes one line:

claude mcp add deep-research -- uv run --directory /path/to/deep-research-mcp deep-research-mcp

Stdio for local spawning, Streamable HTTP on /mcp for a shared server. The repository also ships a CLI (cli/deep-research-cli.py) and a Textual TUI (cli/deep-research-tui.py) that can either run the agent directly or act as MCP clients against the HTTP endpoint – useful for debugging providers without involving the model in the loop.

Clarification

Optional, disabled by default in my setup and probably going away at some point. Three small chat models – a triage model, a clarifier, and an instruction builder – decide whether a query is underspecified, ask follow-up questions, and merge the answers into a longer brief before the expensive provider call. The pattern is adapted from OpenAI’s cookbook. On local endpoints (Ollama, llama-server), clarification needs a model with reasonable structured-output behaviour; very small models trip the triage step.

Notes

Long-running tasks can take hours. Client timeouts matter more than the agent’s own – raise MCP_TOOL_TIMEOUT (Claude Code) or tool_timeout_sec (Codex) and poll research_status rather than blocking a single tool call.
Clarification and instruction building remain OpenAI-compatible even when the main provider is Gemini or DR-Tulu; set clarification_base_url separately.
The repo ships a Claude Code skill and a Codex skill so the assistant knows how to use the tools without re-reading the README.

Code, tests, and setup: github.com/pminervini/deep-research-mcp.

Some Notes on Gradient Estimation

Sun, 29 Mar 2026 00:00:00 +0000

Assume we have a scalar function $f(x)$ of interest, such as a reward we want to maximise or a loss we want to minimise, and that $x$ is drawn from a distribution $p_{\theta}(x)$ parameterised by $\theta$. A natural quantity to study is the expected value of $f$ under this distribution,

\[J(\theta) = \mathbb{E}_{x \sim p_{\theta}(x)}[f(x)],\]

and we would like to optimise $J(\theta)$ with respect to $\theta$. If we want to do this via gradient-based optimisation, we need a way to compute, or at least estimate, the gradient $\nabla_{\theta} J(\theta)$. The difficulty is that $\theta$ enters $J$ through the sampling distribution rather than through the integrand directly, so we cannot simply differentiate $f$ and call it a day. The rest of these notes collects some of the standard tricks for getting around this.

REINFORCE-Based Estimators

REINFORCE

Assuming $f$ depends on $\theta$ only through the sample $x$,

\[\begin{aligned} \nabla_{\theta} J(\theta) &= \nabla_{\theta} \int p_{\theta}(x) f(x)\, dx \\ &= \int f(x)\, \nabla_{\theta} p_{\theta}(x)\, dx \\ &= \int f(x)\, p_{\theta}(x)\, \nabla_{\theta} \log p_{\theta}(x)\, dx \\ &= \mathbb{E}_{x \sim p_{\theta}(x)} \left[ f(x)\, \nabla_{\theta} \log p_{\theta}(x) \right]. \end{aligned}\]

This is the REINFORCE, or score-function, gradient estimator (Williams, 1992).

REINFORCE + Baseline

For any baseline $b$ that does not depend on the sampled variable $x$,

\[\mathbb{E}_{x \sim p_{\theta}(x)} \left[ \left( f(x) - b \right) \nabla_{\theta} \log p_{\theta}(x) \right] = \nabla_{\theta} J(\theta),\]

because

\[\mathbb{E}_{x \sim p_{\theta}(x)} \left[ b \nabla_{\theta} \log p_{\theta}(x) \right] = b \nabla_{\theta} \int p_{\theta}(x)\, dx = 0.\]

So subtracting such a baseline preserves unbiasedness. For a scalar parameter, writing $s(x) = \nabla_{\theta} \log p_{\theta}(x)$, the variance is

\[\operatorname{Var}\left[(f(x) - b)s(x)\right] = \mathbb{E}\left[(f(x) - b)^{2} s(x)^{2}\right] - \left(\nabla_{\theta} J(\theta)\right)^{2}.\]

As a function of $b$, this is a quadratic:

\[\mathbb{E}[s(x)^{2}]\, b^{2} - 2 \mathbb{E}[f(x) s(x)^{2}]\, b + \text{const},\]

so it is minimized at

\[b^{\star} = \frac{\mathbb{E}[f(x) s(x)^{2}]}{\mathbb{E}[s(x)^{2}]}.\]

Therefore an appropriate baseline reduces variance, and the optimal constant baseline is $b^{\star}$.

Proximal Policy Optimization

Let $p_{\mathrm{old}}$ be the sampling policy and let $A(x)$ be an advantage estimate treated as constant with respect to $\theta$. The unclipped policy-gradient surrogate is

\[L_{\mathrm{PG}}(\theta) = \mathbb{E}_{x \sim p_{\mathrm{old}}(x)} \left[ r_{\theta}(x)\, A(x) \right], \qquad r_{\theta}(x) = \frac{p_{\theta}(x)}{p_{\mathrm{old}}(x)}.\]

Then

\[\begin{aligned} \nabla_{\theta} L_{\mathrm{PG}}(\theta) &= \mathbb{E}_{x \sim p_{\mathrm{old}}(x)} \left[ A(x)\, \nabla_{\theta} r_{\theta}(x) \right] \\ &= \mathbb{E}_{x \sim p_{\mathrm{old}}(x)} \left[ A(x)\, r_{\theta}(x)\, \nabla_{\theta} \log p_{\theta}(x) \right]. \end{aligned}\]

So PPO is just an importance-weighted REINFORCE estimator. Its clipped objective replaces $r_{\theta}(x) A(x)$ by

\[\min \left( r_{\theta}(x) A(x), \operatorname{clip}(r_{\theta}(x), 1-\varepsilon, 1+\varepsilon) A(x) \right),\]

which keeps the same REINFORCE form when the ratio is unclipped and truncates the update when it is too large (Schulman et al., 2017).

Group Relative Policy Optimization

For the same prompt $q$, let $o_{1}, \ldots, o_{G} \sim \pi_{\mathrm{old}}(\cdot \mid q)$ be a group of sampled outputs with rewards $R_{1}, \ldots, R_{G}$. GRPO uses the group-normalized advantage

\[\hat{A}_{i} = \frac{R_{i} - \operatorname{mean}(R_{1}, \ldots, R_{G})} {\operatorname{std}(R_{1}, \ldots, R_{G})}.\]

Ignoring clipping and the reference-policy KL term, the surrogate objective is

\[L_{\mathrm{GRPO}}(\theta) = \mathbb{E} \left[ \frac{1}{G} \sum_{i=1}^{G} \frac{1}{|o_{i}|} \sum_{t=1}^{|o_{i}|} r_{i,t}(\theta)\, \hat{A}_{i} \right],\]

with token-level importance ratio

\[r_{i,t}(\theta) = \frac{\pi_{\theta}(o_{i,t} \mid q, o_{i,<t})} {\pi_{\mathrm{old}}(o_{i,t} \mid q, o_{i,<t})}.\]

Therefore

\[\nabla_{\theta} L_{\mathrm{GRPO}}(\theta) = \mathbb{E} \left[ \frac{1}{G} \sum_{i=1}^{G} \frac{1}{|o_{i}|} \sum_{t=1}^{|o_{i}|} \hat{A}_{i}\, r_{i,t}(\theta)\, \nabla_{\theta} \log \pi_{\theta}(o_{i,t} \mid q, o_{i,<t}) \right].\]

So GRPO is again REINFORCE with importance weighting, but the baseline is estimated from the rewards of the other samples in the same group rather than from a critic; the full GRPO objective then adds PPO-style clipping and a KL penalty to a reference policy (Shao et al., 2024).

REINFORCE Leave-One-Out

For $K$ independent samples $x_{1}, \ldots, x_{K} \sim p_{\theta}(x)$, define the leave-one-out baseline

\[b_{k} = \frac{1}{K-1} \sum_{j \neq k} f(x_{j}).\]

Then

\[\widehat{\nabla_{\theta} J} = \frac{1}{K} \sum_{k=1}^{K} \left[ \left( f(x_{k}) - b_{k} \right) \nabla_{\theta} \log p_{\theta}(x_{k}) \right].\]

This is unbiased because, conditional on $x_{-k}$, the baseline $b_{k}$ does not depend on $x_{k}$, so

\[\mathbb{E} \left[ b_{k} \nabla_{\theta} \log p_{\theta}(x_{k}) \mid x_{-k} \right] = b_{k}\, \mathbb{E} \left[ \nabla_{\theta} \log p_{\theta}(x_{k}) \right] = 0.\]

This is the leave-one-out control-variate idea underlying REINFORCE Leave-One-Out (RLOO) and VIMCO (Mnih and Rezende, 2016; Ahmadian et al., 2024).

Augment-REINFORCE-Merge / DisARM

For a Bernoulli random variable $z \sim \mathrm{Bernoulli}(\sigma(\phi))$, the Augment-REINFORCE-Merge (ARM) estimator is

\[\nabla_{\phi} \mathbb{E}[f(z)] = \mathbb{E}_{u \sim \mathrm{Uniform}(0, 1)} \left[ \left( f(\mathbf{1}[u > \sigma(-\phi)]) - f(\mathbf{1}[u < \sigma(\phi)]) \right) \left( u - \frac{1}{2} \right) \right].\]

DisARM integrates out the continuous augmentation and, for Bernoulli logits $\alpha_{\theta}$ and an antithetic pair $(b, \tilde{b})$, yields

\[\widehat{\nabla_{\theta} J} = \sum_{i} \left[ \frac{1}{2} \left( f(b) - f(\tilde{b}) \right) (-1)^{\tilde{b}_{i}} \mathbf{1}[b_{i} \neq \tilde{b}_{i}] \sigma(|(\alpha_{\theta})_{i}|) \nabla_{\theta} (\alpha_{\theta})_{i} \right].\]

These are unbiased binary estimators based on antithetic sampling (Yin and Zhou, 2019; Dong et al., 2020).

REBAR

Let $b = H(z)$ with $z \sim p(z \mid \theta)$ and let $\tilde{z} \sim p(z \mid b, \theta)$ be a conditional relaxed sample. With the control variate $c(z) = \eta f(\sigma_{\lambda}(z))$, REBAR uses the unbiased estimator

\[\widehat{\nabla_{\theta} J} = \left[ f(b) - \eta f(\sigma_{\lambda}(\tilde{z})) \right] \nabla_{\theta} \log p(b \mid \theta) + \eta \nabla_{\theta} f(\sigma_{\lambda}(z)) - \eta \nabla_{\theta} f(\sigma_{\lambda}(\tilde{z})).\]

So REBAR is REINFORCE plus a debiased continuous-relaxation control variate (Tucker et al., 2017).

RELAX

RELAX replaces the hand-designed REBAR control variate by a learned differentiable surrogate $c_{\phi}$. With the same notation as above,

\[\widehat{\nabla_{\theta} J} = \left[ f(b) - c_{\phi}(\tilde{z}) \right] \nabla_{\theta} \log p(b \mid \theta) + \nabla_{\theta} c_{\phi}(z) - \nabla_{\theta} c_{\phi}(\tilde{z}).\]

This estimator is unbiased for any $c_{\phi}$; the point is to learn $c_{\phi}$ so as to reduce variance (Grathwohl et al., 2018).

Pathwise Derivative / Reparameterization Trick

If we can write the random variable as a differentiable transformation $x = g_{\theta}(\varepsilon)$ with $\varepsilon \sim p(\varepsilon)$ independent of $\theta$, then

\[\begin{aligned} \nabla_{\theta} J(\theta) &= \nabla_{\theta} \mathbb{E}_{\varepsilon \sim p(\varepsilon)} \left[ f(g_{\theta}(\varepsilon)) \right] \\ &= \mathbb{E}_{\varepsilon \sim p(\varepsilon)} \left[ \nabla_{\theta} f(g_{\theta}(\varepsilon)) \right] \\ &= \mathbb{E}_{\varepsilon \sim p(\varepsilon)} \left[ \frac{\partial f}{\partial x} \frac{\partial g_{\theta}(\varepsilon)}{\partial \theta} \right]. \end{aligned}\]

This is the pathwise, or reparameterization, gradient estimator (Kingma and Welling, 2013; Rezende et al., 2014).

Evolution Strategies

For a Gaussian search distribution over parameters,

\[J(\theta) = \mathbb{E}_{\varepsilon \sim \mathcal{N}(0, I)}[f(\theta + \sigma \varepsilon)] = \mathbb{E}_{x \sim \mathcal{N}(\theta, \sigma^{2} I)}[f(x)].\]

Applying the same identity,

\[\begin{aligned} \nabla_{\theta} J(\theta) &= \mathbb{E}_{x \sim \mathcal{N}(\theta, \sigma^{2} I)} \left[ f(x)\, \nabla_{\theta} \log \mathcal{N}(x; \theta, \sigma^{2} I) \right] \\ &= \mathbb{E}_{x \sim \mathcal{N}(\theta, \sigma^{2} I)} \left[ f(x)\, \frac{x - \theta}{\sigma^{2}} \right] \\ &= \frac{1}{\sigma}\, \mathbb{E}_{\varepsilon \sim \mathcal{N}(0, I)} \left[ f(\theta + \sigma \varepsilon)\, \varepsilon \right]. \end{aligned}\]

Equivalently, if we define the scaled perturbation $\delta = \sigma \varepsilon$, then $\delta \sim \mathcal{N}(0, \sigma^{2} I)$ and

\[\nabla_{\theta} J(\theta) = \frac{1}{\sigma^{2}}\, \mathbb{E}_{\delta \sim \mathcal{N}(0, \sigma^{2} I)} \left[ f(\theta + \delta)\, \delta \right].\]

This is the same estimator: it is only a change of variables. Indeed, substituting $\delta = \sigma \varepsilon$ gives

\[\frac{1}{\sigma}\, \mathbb{E}_{\varepsilon \sim \mathcal{N}(0, I)} \left[ f(\theta + \sigma \varepsilon)\, \varepsilon \right] = \frac{1}{\sigma^{2}}\, \mathbb{E}_{\delta \sim \mathcal{N}(0, \sigma^{2} I)} \left[ f(\theta + \delta)\, \delta \right].\]

So Evolution Strategies are just REINFORCE applied to a Gaussian distribution over parameters; the standardized-noise form $\varepsilon \sim \mathcal{N}(0, I)$ is the one used by (Salimans et al., 2017), while the scaled-perturbation form $\delta \sim \mathcal{N}(0, \sigma^{2} I)$ is the notation used by (Huszár, 2017) and (Barber, 2017).

Straight-Through Estimation

Let $z = H(a_{\theta}(\xi))$ be a hard discrete decision, where $H$ is a step function or an $\arg\max$. The exact pathwise derivative is unusable because $\partial H / \partial a$ is zero almost everywhere or undefined. The Straight-Through Estimator (STE) replaces the backward derivative by that of a smooth surrogate $\tilde{H}$:

\[\frac{\partial z}{\partial a} \; := \; \frac{\partial \tilde{H}(a)}{\partial a}.\]

Therefore, for $J(\theta) = \mathbb{E}_{\xi}[L(z)]$,

\[\nabla_{\theta} J(\theta) \approx \mathbb{E}_{\xi} \left[ \frac{\partial L}{\partial z} \frac{\partial \tilde{H}(a_{\theta}(\xi))}{\partial a} \frac{\partial a_{\theta}(\xi)}{\partial \theta} \right].\]

This is not an exact identity: STE is a biased surrogate-gradient estimator (Bengio et al., 2013).

Implicit Maximum Likelihood Estimation

Let $\operatorname{MAP}(\theta)$ denote the maximum a posteriori solution, i.e. the feasible discrete structure that maximizes the score induced by $\theta$. In practice this can be top-$k$ selection, shortest-path inference via Dijkstra, or an integer linear programming solver. If we add random perturbations $\varepsilon \sim \rho$ to the scores and solve the perturbed optimization problem, we obtain

\[z_{\theta}(\varepsilon) = \operatorname{MAP}(\theta + \varepsilon), \qquad \varepsilon \sim \rho.\]

The map $\varepsilon \mapsto \operatorname{MAP}(\theta + \varepsilon)$ pushes the perturbation law $\rho$ forward to an implicit distribution over feasible solutions, $p_{\mathrm{PM}}(z; \theta) = \mathbb{P}_{\varepsilon \sim \rho} \left[ z = \operatorname{MAP}(\theta + \varepsilon) \right].$ Sampling from this implicit distribution by drawing $\varepsilon \sim \rho$ and solving one perturbed MAP problem is the perturb-and-MAP construction introduced by (Papandreou and Yuille, 2011).

Let $g = \nabla_{z} f(z_{\theta}(\varepsilon))$ be the downstream gradient and define target parameters

\[\theta^\prime = \theta - \lambda g.\]

Here $\theta^\prime$ is chosen so as to move in a direction expected to reduce the downstream loss. This is the key idea behind perturbation-based implicit differentiation: in the smooth-marginal setting, perturbing parameters in the negative downstream-gradient direction yields a lower-loss solution (Domke, 2010), and I-MLE adapts this idea to the discrete perturb-and-MAP setting.

For a constrained exponential-family distribution

\[p(z; \theta) \propto \exp(\langle z, \theta \rangle - A(\theta)),\]

the negative log-likelihood of one observation $\hat{z}$ is

\[\ell(\theta; \hat{z}) = A(\theta) - \langle \hat{z}, \theta \rangle,\]

so its gradient is

\[\nabla_{\theta} \ell(\theta; \hat{z}) = \nabla_{\theta} A(\theta) - \hat{z} = \mu(\theta) - \hat{z},\]

where $\mu(\theta) = \mathbb{E}_{p(z; \theta)}[z]$. More generally, if the target is itself a distribution $q(z; \theta^\prime)$ from the same exponential family, the implicit maximum-likelihood objective is

\[L(\theta, \theta^\prime) = -\mathbb{E}_{z \sim q(z; \theta^\prime)}[\log p(z; \theta)],\]

and its gradient is

\[\nabla_{\theta} L(\theta, \theta^\prime) = \mu(\theta) - \mathbb{E}_{z \sim q(z; \theta^\prime)}[z] = \mu(\theta) - \mu(\theta^\prime).\]

Implicit Maximum Likelihood Estimation (I-MLE) chooses the target distribution $q$ through $\theta^\prime$ and then approximates these marginals by perturb-and-MAP samples:

\[\widehat{\nabla_{\theta} J} \propto z_{\theta}(\varepsilon) - z_{\theta^\prime}(\varepsilon) = \operatorname{MAP}(\theta + \varepsilon) - \operatorname{MAP}(\theta - \lambda g + \varepsilon).\]

So the I-MLE estimator is obtained by taking the maximum-likelihood gradient $\mu(\theta) - \mu(\theta^\prime)$ for exponential-family distributions and replacing the intractable marginals by perturb-and-MAP approximations $\mu(\theta) \approx \mathbb{E}_{\varepsilon}[z_{\theta}(\varepsilon)], \qquad \mu(\theta^\prime) \approx \mathbb{E}_{\varepsilon}[z_{\theta^\prime}(\varepsilon)].$

With multiple independent perturbations, one uses

\[\widehat{\nabla_{\theta} J} = \frac{1}{S} \sum_{i=1}^{S} \left[ \operatorname{MAP}(\theta + \varepsilon_{i}) - \operatorname{MAP}(\theta - \lambda g_{i} + \varepsilon_{i}) \right], \qquad \varepsilon_{i} \sim \rho, \qquad g_{i} = \nabla_{z} f(z_{\theta}(\varepsilon_{i})),\]

with the factor $1 / \lambda$ often absorbed into the learning rate. This is a biased finite-difference estimator based on an implicit target distribution (Niepert et al., 2021).

Adaptive I-MLE (AIMLE) extends this by adapting the perturbation scale $\lambda$ during training, explicitly trading off sparsity and bias in the estimate, and by also considering centered finite differences (Minervini et al., 2022).

Gumbel-Softmax / Gumbel-Sigmoid

For a categorical variable with logits $\ell_{\theta}$, sample i.i.d. Gumbel noise $g_{i} = - \log(-\log u_{i})$ with $u_{i} \sim \mathrm{Uniform}(0, 1)$, and define

\[y_{i} = \frac{\exp((\ell_{\theta, i} + g_{i}) / \tau)} {\sum_{j} \exp((\ell_{\theta, j} + g_{j}) / \tau)}.\]

Then $y$ is a differentiable function of $\theta$ and noise $g$, so the pathwise gradient applies:

\[\nabla_{\theta} J(\theta) = \mathbb{E}_{g} \left[ \nabla_{\theta} f(y_{\theta}(g)) \right] = \mathbb{E}_{g} \left[ \frac{\partial f}{\partial y} \frac{\partial y_{\theta}(g)}{\partial \theta} \right].\]

This is the pathwise, or reparameterization, gradient estimator underlying Gumbel-Softmax (Jang et al., 2017; Maddison et al., 2017).

In the binary case, often called Gumbel-Sigmoid or Binary Concrete, sample logistic noise $\lambda = \log u - \log(1-u)$ with $u \sim \mathrm{Uniform}(0, 1)$ and define

\[y = \sigma \left( \frac{a_{\theta} + \lambda}{\tau} \right).\]

Then the same reparameterization argument gives

\[\nabla_{\theta} \mathbb{E}_{\lambda}[f(y_{\theta}(\lambda))] = \mathbb{E}_{\lambda} \left[ \frac{\partial f}{\partial y} \frac{\partial y_{\theta}(\lambda)}{\partial \theta} \right].\]

Measure-Valued Derivatives

Suppose the derivative of the measure can be decomposed as

\[\nabla_{\theta_{i}} p(x; \theta) = c_{\theta_{i}} \left[ p_{i}^{+}(x; \theta) - p_{i}^{-}(x; \theta) \right].\]

Then

\[\begin{aligned} \nabla_{\theta_{i}} J(\theta) &= \int f(x)\, \nabla_{\theta_{i}} p(x; \theta)\, dx \\ &= c_{\theta_{i}} \left( \mathbb{E}_{p_{i}^{+}(x; \theta)}[f(x)] - \mathbb{E}_{p_{i}^{-}(x; \theta)}[f(x)] \right). \end{aligned}\]

So the gradient becomes a scaled difference of two expectations under the positive and negative parts of the weak derivative (Mohamed et al., 2020).

Real-World Impact of Our Research

Tue, 15 Jul 2025 01:00:00 +0100

Academic research sometimes risks to be disconnected from real-world applications. Research from our group demonstrated significant real-world impact across multiple domains, from improving the efficiency of LLM inference and training to new state-of-the-art evaluation protocols, and contributing to several industry products.

Industry Adoption

KV Cache Compression

Our work on KV cache compression, presented at EMNLP 2024 as an oral presentation (top 9% of papers), has been directly adopted by NVIDIA’s KV Press library. This library is now widely used across the industry for reducing the inference footprint of LLMs

MMLU-Redux

Our MMLU-Redux benchmark, presented at NAACL 2025, is being widely adopted for evaluating LLMs across multiple frontier labs. After identifying a concerning 57% error rate in the original MMLU benchmark’s Virology subset, we created a manually curated, expert-verified dataset of 5,700 questions. Some industry adopters include DeepSeek, Alibaba Qwen, Tencent Hunyuan, Moonshot KIMI, NVIDIA NeMo Skills, and LG AI Research EXAONE.

Frontier Model Pre-Training

One of our pre-training techniques from “Analysing The Impact of Sequence Composition on Language Model Pre-Training” (a project led by Yu Zhao and presented at ACL 2024), specifically intra-document causal masking, was adopted by Meta in training Llama3, their flagship line of LLMs, by Hugging Face for training their SmolLM3 models, and by Allen AI for training their OLMo 3 models.

Complex Query Answering

Our work on answering complex queries on large and incomplete Knowledge Graphs (“Complex Query Answering with Neural Link Predictors”, which received an Outstanding Paper Award at ICLR 2021), is also having some impact beyond academia. Based on personal communications, our complex query answering techniques are being adopted by several tech companies to develop their AI products (I am happy to discuss this verbally). We are now continuing this line of research – for example, see our NeurIPS 2023 and ICML 2025 papers on this topic.

MMLongBench

MMLongBench, a benchmark for evaluating multimodal long-context models by Zhaowei Wang et al., is being adopted by Xiaomi in their MiMo-VL project, a vision-language model for general visual understanding and multimodal reasoning.

Media Coverage

Our research received some media coverage. For example, our work “Lost in Time: Clock and Calendar Understanding Challenges in Multimodal LLMs”, led by the amazing Rohit Saxena, received extensive media coverage, including The New York Times, VICE, Gizmodo, The Engineer, Yahoo! News, and Tech Xplore.

Our project “Inverse Scaling in Test-Time Compute” led by Aryo Gema in collaboration with Anthropic, also received some media attention, with coverage in VentureBeat, YourStory, AI Times, TechZine, The Interview Times, AI News, MarkTechPost, and others.

March 2025 in Research

Thu, 17 Apr 2025 01:00:00 +0100

We have been working on language model evaluation, knowledge utilization, efficiency, and multimodal reasoning. We had papers at ICLR 2025, NAACL 2025 (x3), AAAI 2025, and others, along with several ongoing works.

NAACL 2025 – Controlling Knowledge & Reasoning

Steering Knowledge Selection Behaviours in LLMs via SAE-Based Representation Engineering, by Yu Zhao et al. – We introduce SpARE, a training‑free method to control whether an LLM relies on its internal parametric knowledge or given context when conflicts arise. By analyzing mid‑layer activations with sparse autoencoders, SpARE identifies conflict signals and manipulates them to steer the model at inference time, significantly improving performance on open‑domain QA compared to prior methods. (Oral presentation)
Are We Done with MMLU?, by Aryo Gema and many others – We analyze the Massive Multitask Language Understanding benchmark, uncovering a fairly high error rate – for example, in the Virology subset, 57% of sampled questions had issues. We introduce MMLU‑Redux, a manually curated subset of 5,700 expert‑verified questions, and show that corrected evaluations can substantially alter model rankings. MMLU‑Redux is open‑sourced also adopted for example by DeepSeek and Qwen!
Self-Training Large Language Models for Tool-Use Without Demonstrations, based on Ne Luo’s MSc project – We explore whether LLMs can learn tool usage (e.g., search engines, calculators) without hand‑crafted examples. Starting with zero‑shot prompts, we generate synthetic tool‑using traces and then fine‑tune the model with them. On PopQA, the self‑trained model gains +3.7% accuracy, though results vary on other datasets, highlighting both promise and challenges in autonomous tool‑use learning. Ne Luo is looking for a PhD position, contact her if you are interested in working with her!

ICLR 2025 – Learning & Evaluation

An Auditing Test to Detect Behavioral Shift in Language Models, by the amazing Leo Richter – We propose a method for continual Behavioral Shift Auditing (BSA) of LLMs. This statistical test monitors an LLM’s outputs for significant deviations from a reference model’s behavior, with theoretical guarantees on detecting genuine shifts while avoiding false alarms. Our BSA approach relies on catching subtle changes in a model’s toxicity and translation performance after fine-tuning, using only a few hundred examples, offering a practical tool to ensure that an LLM remains aligned during its deployment/lifetime.

Reasoning and Planning for Large Language Models

Lost in Time: Clock and Calendar Understanding Challenges in Multimodal LLMs, by Rohit Saxena et al. – We introduce ClockQA and CalendarQA for testing multimodal LLMs’ temporal reasoning from images, revealing widespread failures and motivating models with better time‑date understanding. This article got plenty of media coverage – e.g., on Gizmodo, The Engineer, Yahoo! News, VICE, Tech Xplore, and many others!

AAAI 2025 – Efficient Inference

Adaptive Computation Modules: Granular Conditional Computation for Efficient Inference, by Bartosz Wójcik, Alessio Devoto, et al. – We propose Adaptive Computation Modules (ACMs) for dynamic, per‑token computation in Transformers. ACMs consist of cascaded sub‑modules with gating functions that allow easy tokens to exit early. We propose a distillation method to retrofit pre‑trained models with ACMs, cutting inference cost without accuracy loss in vision and speech tasks, offering a plug‑and‑play approach to more efficient AI systems!

COLING 2025 – Multilingual Resources

SynDARin: Synthesising Datasets for Automated Reasoning in Low-Resource Languages, by Gayane Ghazaryan, Erik Arakelyan (who just got his PhD and joined NVIDIA! 🚀) et al. – SynDARin synthesizes QA datasets in low‑resource languages (e.g., Armenian) by generating English questions via LLMs from parallel corpora, translating and validating them. The resulting 16,000+ QA pairs produce a challenging benchmark where models often perform near chance, highlighting critical gaps and enabling rapid evaluation in languages lacking resources.

Frontiers in AI 2025 – Human-AI Collaboration

Fostering Effective Hybrid Human-LLM Reasoning and Decision Making – We examine frameworks combining LLMs and human judgment for complex tasks, offering design principles for AI‑assisted decision systems. Through case studies, we show that integrating LLM‑generated insights with human oversight yields more reliable and interpretable outcomes than either alone, providing guidelines for principled human‑in‑the‑loop systems.

What’s Brewing

Noiser: Bounded Input Perturbations for Attributing Large Language Models, by Reza Madani et al. – Noiser perturbs input embeddings to attribute token importance, introducing an “answerability” check to validate attributions. Outperforming gradients and attention, Noiser offers robust post‑hoc explanations for LLM predictions.
An Analysis of Decoding Methods for LLM-based Agents for Faithful Multi-Hop Question Answering, by Alex, Sanad, and others amazing students at the UoE – We analyse how faithfulness‑enhancing decoding (e.g., DeCoRe) within the ReAct agent framework improves multi‑hop QA, boosting HotpotQA F1 from 19.5 to 32.6, underscoring the role of decoding in reliable LLM reasoning.
Q-Filters: Leveraging QK Geometry for Efficient KV Cache Compression, led by Nathan Godey – Q-Filters uses query‑key geometric projections to filter past tokens on the fly, compressing KV cache without retraining and matching attention‑based methods like SnapKV, enabling long‑context generation with minimal memory.
PosterSum: A Multimodal Benchmark for Scientific Poster Summarization, by the amazing Rohit Saxena – PosterSum offers 16,000+ re search posters paired with abstracts for evaluating vision‑language summarization. Our “Segment & Summarize” approach secures a 3.1% ROUGE‑L gain, highlighting this benchmark’s challenge.

Postdoc Position in Multimodal Foundation Models

Sat, 01 Mar 2025 00:00:00 +0000

Amazing opportunity to join our team at the School of Informatics, University of Edinburgh! The School of Informatics is seeking a Postdoctoral Research Associate to work on evaluating and improving multimodal foundation models, with a particular focus on Vision-Language Models (VLMs).

About the Position

This is a full-time position running until January 2029, fully funded by the AI Hub in Generative Models. The successful candidate will join the Edinburgh NLP Group, one of the best NLP research groups in the world!

Key Details:

Duration: Fixed-term contract until January 2029
Application Deadline: April 8th, 2025, 12:59 AM (UK time)

For more details or to apply, visit https://edin.ac/3DDQK1o

For informal enquiries, feel free to reach out to me directly at p.minervini@ed.ac.uk

November 2024 in Research

Fri, 01 Nov 2024 00:00:00 +0000

My amazing collaborators will be presenting three papers at EMNLP 2024 (main track), a leading conference in natural language processing, happening in Miami later this month! A few weeks ago I also blogged about our ACL 2024, ICML 2024, and CoLM 2024 papers – you can check the post here.

Our work at EMNLP 2024

We will be presenting three papers this year at EMNLP, a flagship NLP conference:

A Simple and Effective $L_{2}$ Norm-Based Strategy for KV Cache Compression, by Yu Zhao, Alessio Devoto, et al. – we introduce a simple strategy for compressing the Key-Value (KV) cache in large language models by utilizing the $L_{2}$ norm of key embeddings; specifically, we found a correlation between low $L_{2}$ norms and high attention scores, allowing them to identify influential KV pairs before querying. For example, here we can see the attention distributions for five heads at layer 9 in Llama2-7B – we can see that the attention scores (top) and the key $L_{2}$ norms (bottom) are highly correlated. — Our method effectively reduces KV cache size by up to 90% without loss of accuracy and is compatible with FlashAttention. This paper will be presented as an Oral – top 8% of the accepted papers! An extended version of this paper will also be presented at the Efficient Natural Language and Speech Processing workshop at NeurIPS 2024! EMNLP Poster
Atomic Inference for NLI with Generated Facts as Atoms, by Joe Stacey et al. – we propose an atomic inference approach for Natural Language Inference (NLI) that decomposes inputs into individual facts or atoms, and explicitly models the entailment relationships between such atoms. Furthermore, we propose a multi-stage fact generation process and a specialized training regime for incorporates such facts, achieving state-of-the-art results in several hard NLI tasks. Our best system, FGLR, produces significantly more robust and accurate results than large-scale language models while providing clear interpretability guarantees by identifying the specific atoms responsible for each prediction! Joe wrote an amazing blog post on this work, check it out!
Unveiling and Consulting Core Experts in Retrieval-Augmented MoE-based LLMs, by Xin Zhou, Ping Nie et al. – we analyse Mixture-of-Expert (MoE)-based Large Language Models (LLMs) in the context of Retrieval-Augmented Generation (RAG); we identify the groups of experts that are primarily responsible for RAG-related behaviors, such as identifying whether the parametric knowledge is sufficient to solve a given knowledge-intensive task; assessing the quality of retrieved documents; and improving the utilisation of context. Based on these findings, we propose several strategies to improve the efficiency and effectiveness of RAG systems by adjusting expert activations.

What’s brewing

We have several super-interesting works in the pipeline! Here are some of them:

Steering Knowledge Selection Behaviours in LLMs via SAE-Based Representation Engineering, by Yu Zhao et al. – we introduce SpARE, a training-free method that leverages pre-trained sparse auto-encoders (SAEs) to control the knowledge selection behavior of large language models (LLMs) when faced with conflicts between their internal (parametric) knowledge and external (contextual) information. By identifying and manipulating functional features within the LLMs’ internal activations, SpARE can steer the model to prioritize either parametric or contextual knowledge during inference. We show that SpARE is surprisingly effective at resolving knowledge conflicts in open-domain question-answering tasks, producing significantly better results than existing representation engineering and contrastive decoding methods. The insights in this paper are based on another paper, Analysing the Residual Stream of Language Models Under Knowledge Conflicts also by Yu Zhao et al. that will appear in the Workshop on Foundation Model Interventions @ NeurIPS 2024.
Mixtures of In-Context Learners, by Giwon Hong et al. – we propose Mixtures of In-Context Learners (MoICL), a method that trains a set of experts via in-context learning, and learns a weighting function to merge their outputs, addressing many of the limitations of standard in-context learning (ICL). MoICL yields significantly more accurate results than many strong baselines (up to +13% compared to ICL and LENS); reduces inference time by achieving similar performance with fewer demonstrations; and shows greater robustness to out-of-domain, imbalanced, or noisy demonstrations.
DeCoRe: Decoding by Contrasting Retrieval Heads to Mitigate Hallucinations, by Aryo Gema et al. – we introduce DeCoRe (Decoding by Contrasting Retrieval Heads), a novel, training-free decoding strategy designed to mitigate hallucinations in large language models (LLMs). DeCoRe works by masking specific retrieval heads — attention heads responsible for extracting relevant contextual information — to induce hallucinations, and then contrast the outputs of the base LLM and the masked LLM, using conditional entropy as a guide. DeCoRe significantly improves performance on tasks requiring high contextual faithfulness, such as summarization, instruction following, and open-book question answering, and surprisingly (to us), it also helps with factual recall!
FLARE: Faithful Logic-Aided Reasoning and Exploration, by Erik Arakelyan et al. – we introduce FLARE (Faithful Logic-Aided Reasoning and Exploration), a framework designed to improve the reasoning abilities of LLMs in knowledge-intensive reasoning tasks. FLARE use an intermediate logic programming-inspired representation of the reasoning process by generating Prolog code and simulating a program execution, ensuring that the reasoning process remains faithful and interpretable without relying on external solvers. FLARE achieves state-of-the-art results on seven out of nine diverse reasoning benchmarks, and we identify a strong correlation between the faithfulness of the reasoning process and the downstream model accuracy.

July 2024 in Research

Mon, 01 Jul 2024 01:00:00 +0100

My amazing collaborators will be presenting several works at ACL 2024, ICML 2024, and CoLM 2024 in the upcoming weeks/months!

Our work at ACL 2024

We will be presenting four papers this year at ACL, the flagship NLP conference:

Analysing The Impact of Sequence Composition on Language Model Pre-Training, by Yu Zhao et al. – we analyse several language model pre-training schemes and find out that, e.g., intra-document causal masking helps both in terms of pre-training dynamics, and accuracy on a wide array of downstream tasks! This approach was later adopted by Llama 3, Meta’s flagship language model family. This paper will be presented as an Oral – top 8% of the accepted papers!
SparseFit: Few-shot Prompting with Sparse Fine-tuning for Jointly Generating Predictions and Natural Language Explanations, by Jesus Solano, Mardhiyah Sanni et al. – we introduce SparseFit, a sparse fine-tuning method that uses few-shot prompting and discrete prompts to efficiently generate both predictions and natural language explanations (NLEs) with large pre-trained language models, achieving competitive performance while significantly reducing the number of fine-tuning parameters! [poster]
Using Natural Language Explanations to Improve Robustness of In-context Learning, by Xuanli He et al. – we found that integrating NLEs into in-context learning significantly improves the robustness of large language models against adversarial inputs, and show that generating NLEs with frontier models in a few-shot setting can significantly improve accuracy on challenging natural language inference tasks compared to traditional in-context learning and human-generated NLEs.
Probing the Emergence of Cross-lingual Alignment during LLM Training, by Hetong Wang et al. – we analyse how cross-lingual alignment emerges during the training of multilingual large language models by probing neuron activity in different languages. We find that higher neuron overlap between languages correlates strongly with improved zero-shot cross-lingual transfer performance, but also identifies phases during training where both alignment and performance degrade, offering new insights into the dynamics of multilingual model training! [poster]

Our work at ICML 2024

We will be presenting three works this year at ICML – one in the main conference and two in co-located workshops:

On the Independence Assumption in Neurosymbolic Learning, by Emile van Krieken et al. – we analyse the common assumption in neurosymbolic learning that symbols are conditionally independent given the input, and argue that this assumption biases models towards deterministic solutions and limits their ability to express uncertainty.
Attention Is All You Need But You Don’t Need All Of It For Inference of Large Language Models, by Georgy Tyukin et al. – we investigate the effects of removing MLP and attention layers in large language models during inference, finding that removing deeper attention layers can only marginally reduce performance while significantly improving inference speed!
An Auditing Test to Detect Behavioral Shift in Language Models, by Leo Richter et al. – we propose a continuous online auditing framework to detect behavioural shifts in language models, ensuring that deployed models remain aligned with societal values and preventing vendors or attackers from covertly deploying unaligned models for malicious purposes.

Our work at CoLM 2024

The Conference on Language Modeling (CoLM) is a very new thing. I have been area-chairing for CoLM this year, and I’m really impressed by the quality of all submissions! We will be presenting two papers:

Forklift: An Extensible Neural Lifter, by Jordi Armengol-Estapé et al. – we introduce Forklift, a framework that uses neural models to translate assembly code across different instruction set architectures by “lifting” source assembly code into an intermediate representation, thereby reducing the engineering effort required for cross-architecture software migration!
Evaluating the Adversarial Robustness of Retrieval-Based In-Context Learning for Large Language Models, by Simon Yu, Jie He et al. – we analyse the adversarial robustness of retrieval-based in-context learning, finding that while retrieval-augmented methods improve robustness against test sample attacks, they increase vulnerability to adversarially perturbed demonstrations; to address this, we propose a new training-free defence method, which significantly improves adversarial robustness.

Looking for Postdocs, June 2024 Edition

Sat, 01 Jun 2024 01:00:00 +0100

We have an opening for a 3-year postdoc – more details are available here – on a project funded by Huawei via the Huawei-Edinburgh Joint Lab initiative, with me as the Principal Investigator (PI).

The researcher will work on projects involving the design and application of improving the robustness and trustworthiness of Large Language Models when solving complex reasoning tasks, while improving their explainability and generalisation properties. They will be part of the Edinburgh NLP Group, a world-leading research group in Natural Language Processing.

Looking for Postdocs!

Tue, 01 Nov 2022 00:00:00 +0000

We have an opening for a 2-year postdoc – more details are available here – on a project titled Gradient-based Learning of Complex Latent Structures, with me as the Principal Investigator (PI), and Antonio Vergari (IANC) and Edoardo Ponti (ILCC) as co-PIs. The position is entirely funded by the Edinburgh Laboratory for Integrated Artificial Intelligence (ELIAI) – if you want to know more, feel free to reach out!

You can apply at this link.

Project description

Imposing structural constraints on the latent representations learned by deep neural models has several applications, which can improve their explainability, their robustness, and their ability to generalise to out-of-domain distributions. For example, we can learn more explainable models by making them selectively decide which parts of the input to consider; and we can improve their generalisation properties by learning representations suitable for reasoning tasks, such as deductive reasoning and planning, and comply with any desired constraints. For instance, the intermediate structure can represent a relational graph between objects in the world; the relationships between multiple sub-questions in a complex question; or computation graphs which can be executed to produce a prediction.

In this project, we aim to investigate how we can derive better methods for back-propagating through mixed continuous-discrete complex latent structures, and how we can leverage them for learning more explainable, data-efficient, and robust deep neural models. The reason why discrete latent representations are not widely adopted by deep neural models is that they tend to not interact well with gradient-based optimisation methods, but this started to change recently (e.g., see Niepert et al., 2021; Minervini et al. 2022), enabling a wide range of applications and use cases.

Position

The post holder will work on projects involving the design and application of deep learning models with discrete latent structures for improving their explainability, generalisation, and robustness properties. They will be part of the new Edinburgh Laboratory for Integrated Artificial Intelligence and the Edinburgh NLP Group, a world-leading research group in Natural Language Processing.

The School of Informatics is one of the largest research centres in Computer Science in Europe, and it has been ranked #1 in the UK in terms of research power by a large margin. The Edinburgh NLP Group is consistently ranked among the world’s leading research groups in Natural Language Processing. We are offering an exciting opportunity to work in an interdisciplinary, collaborative, friendly, and supportive environment, integrating different sub-fields of Computer Science and Artificial Intelligence.