Improving Prompt Injection Robustness

Prompt injection persists because models infer who is speaking from writing style rather than from role tags. Ye, Cui & Hadfield-Menell (ICML 2026, https://arxiv.org/abs/2603.12277) show that user-style text wrapped in tool tags is still perceived as coming from the user. Models can resist prompt injection by perceiving which channel text arrived on or memorising what attacks look like, but only the first generalises.

We want to find architecture-level changes that make role information a signal to the model, and to show mechanistically that those changes cause the improvement rather than coinciding with it. The work is hands-on: modify architectures, fine-tune, evaluate both behaviour and internal representations.

We have a few starting directions, and are very open to mentees proposing others.

First, embeddings. ISE (Wu et al., ICLR 2025, https://arxiv.org/abs/2410.09102) and PFT (Wang et al., ICML 2025, https://arxiv.org/abs/2505.00626) use positional encodings to guide role perception. We want to build on their work (e.g. combining them, trying different embedding methods) and evaluating them representationally (going beyond the reported attack success rates to activation level analysis).

Second, attention sinks. Models treat the BOS token specially (Xiao et al. 2023, https://arxiv.org/abs/2309.17453), which may be why system-prompt tokens near the start are easy to mark, and why positional interventions outperform delimiter-based ones. We want to test if this is true and If sinks matter, train model organisms that use more of them.

Prompt injection is one of the cheapest methods for jailbreaking a (frontier) model, and it is surprisingly effective. It requires no gradient access, no fine-tuning, no compute, and no expertise - text on a web page, in a document, or in an API response is enough. As models are given increasingly more tools the surface area for prompt injection expands. That asymmetry, between how little an attack costs and how much control it can obtain, is what makes this a crucial safety problem.

The user role is the channel through which humans authorise consequential actions. A model that cannot reliably distinguish user instructions from third party text can start pursuing unintended goals. A model that cannot distinguish user instructions from its own output (e.g. agent made plans) can effectively manufacture its own authorisation. The downstream effects of these failures get higher as AI becomes more agentic and capable of long horizon tasks.

Mentees own the empirical work. I will help guide scoping, prioritisation, and framing.

High autonomy: implementing and validating the probe and ablation pipelines, designing and running the interventions, debugging the training and evaluation stacks, choosing which diagnostic to run next, and selecting which of the secondary directions to pursue.

Collaborative: interpreting results, judging whether a negative result is real, forming hypothesis and suggesting research directions. These are where research taste is actually learned and I want mentees to develop this skill.

Team size 1–2.

• Solid working knowledge of transformer architecture. You should be comfortable explaining where positional information enters a forward pass, the difference between additive positional embeddings and RoPE, what the residual stream is, and how attention masks work.
• Basic LLM fine-tuning experience (e.g. having trained a small LoRA, familiar with the debugging process, etc.)
• Familiar with AI coding agent use
• Able to run multi-GPU jobs, or willing to learn quickly.
• Understanding of basic interp methods like linear probes and activation steering (or willingness to learn quickly)