kg.

Improving reasoning and safety of large language models

Khashayar
Gatmiry.

Hi there! I am a postdoctoral researcher at the Simons Institute at UC Berkeley, hosted and advised by Peter Bartlett, where I also work closely with Jason Lee and Nika Haghtalab. I am also a part-time researcher on Apple’s Machine Learning Research (MLR) team, working with Emmanuel Abbe and Samy Bengio.

My research explores new training paradigms to advance the reasoning capabilities of large language models while making them more interpretable and safer.

Previously, I completed my PhD in EECS at MIT, where I was fortunate to be advised by Stefanie Jegelka and Jonathan Adam Kelner. My doctoral research focused on optimization for modern deep learning, diffusion models, and fundamental problems in machine learning theory.

Berkeley email ↗Apple email ↗Google Scholar ↗GitHub ↗

News

Recent work

Blackwell-Guided Logit Search (BLS): Optimizing Multi-Turn Jailbreaks with Subliminal Steering

Khashayar Gatmiry, Timothy Ng, Arhaan Garg, Avrajit Ghosh, Emmanuel Abbe, Peter Bartlett, Bin Yu

Preprint · Submitted to ICLR 2027
Blackwell-Guided Logit Search overview: candidate sentences, target-model logit probes, adaptive selection across objectives, and conversation construction.
View full size ↗

PAPERS & PREPRINTS

Publications

23 papers · organized by research topic

Preprint

On the Role of Depth and Looping for In-Context Learning with Task Diversity

Khashayar Gatmiry , Zhiyuan Li, Sashank J. Reddi, Stefanie Jegelka

TransformersPreprint
More on Google Scholar

PREVIOUSLY

A little background

I have spent time at Google in New York (October 2022–January 2023), working with Sashank Reddi and Sobhan Miryouseffi on efficient BERT training, and at Microsoft Research (June–August 2024), working with Ronen Eldan, Adil Salim, and Yi Zhang on language models and embeddings.

ABSTRACT TOKEN CURRICULUM

Learn Your Own Thoughts: Abstract Token Curriculum

Abstract Token Curriculum architecture showing abstract reasoning tokens, hidden embeddings, and the final prediction head.

Teaching a language model to reason often involves showing it worked-out solutions. We explore a different approach: give the model easier problems first, gradually make them harder, and let it develop its own internal scratchpad. We call this Abstract Token Curriculum (ATC). The model learns from the correct answers, without being told what its intermediate steps should be.

On a task that asks whether two nodes in a graph are connected, our method reaches 99.9% accuracy. It reaches 95% accuracy in 51 training epochs, compared with 132 for a competing method that also receives guidance on intermediate steps. In our eight-digit addition experiments, our best variant gets the entire answer right 99.6% of the time, while that competing method scores 0% in the same setup.

The combination matters: on the graph task, giving the model only an easier-to-harder training schedule, or only an internal scratchpad, leaves it performing near chance.

Read paper ↗

HOW TRANSFORMERS LEARN TO SORT

Transformers Learn Two Competing Sorting Algorithms, and Only One Length-Generalizes

Illustrative attention pattern: layer one favors 5, but layer two selects the correct successor, 3.

Two transformers can learn to sort the same training examples almost perfectly—and still learn different algorithms. In our experiments, models trained only on lists of one fixed length settle into two distinct solutions. One effectively sorts with a single attention layer and struggles with longer lists. The other uses both layers and generalizes perfectly to the longer lists we test. Changing the initial weights can steer which solution the model learns.

Longer lists bring a subtle challenge: nearby numbers become more common. Cases that were rare during training can suddenly matter a lot. The two-layer solution divides the work between its layers, helping it handle these cases.

The illustration shows how the layers can cooperate. The first gathers a few candidates—even favoring 5 over the correct answer, 3. That information gives the second layer a narrower range to search, allowing it to focus on 3. Getting the training answers right, it turns out, can hide very different ways of solving the problem.

Read paper ↗

LEARNING FROM SELF-GENERATED FEEDBACK

RLTL;DR: Self-improvement by Internalizing Self-generated Feedback

RLTL;DR training loop: failed attempts produce insights that guide later attempts and are gradually internalized by the agent.

Rewarding successful attempts works well when an AI agent occasionally gets things right. But on a problem it never solves, there is little success to learn from. Our approach, RLTL;DR, gives failed attempts a more useful role: the agent reads the feedback from a checker, writes a short lesson for itself, and tries again with that lesson in mind. Training then helps it absorb those lessons so it can use them without needing the notes.

We test this on coding and tool-use problems where the starting model failed all 128 initial attempts. Standard reinforcement learning leaves first-attempt success at 0–1%. Our method raises it to 12–13%, even with no notes provided at test time.

Surprisingly, training on just 4,000 problem-and-insight pairs recovers almost the full benefit of the larger training process. A brief lesson about what to keep in mind can carry much of the useful learning from an entire attempt.

Read paper ↗

LEARNING TO NAVIGATE LIVE APPS

Online Dense Reinforcement Learning with Planning-Grounding Agent for UI Navigation on Live iOS Apps

UINavBench comparison: UI-TARS 11.0%, Qwen3.8 19.6%, Gemini 3.5 Flash 28.4%, Gemini 3.1 Pro 32.4%, SFT 16.7%, terminal-only PPO 17.6%, PPO 25.5%, and GiGPPO 30.4%.

Given a task such as adding an item to a shopping cart, our 8B planner reads the screenshot, task, and action history, then proposes the next step and updates its remaining plan. A separate, fixed grounding model converts that instruction into an executable action and screen coordinates. Online reinforcement learning updates the planner through interaction with live iOS apps.

What earns a reward? A judge compares before-and-after screenshots against the task and proposed action: progress earns +1, failure earns 0, and acting after completion earns −2. A separate terminal judge rewards task completion. PPO learns from the resulting step-level returns, halving the weight of each successive future dense reward while keeping the terminal reward undiscounted. This limits incentives to prolong a task just to collect intermediate rewards.

What does GiGPPO change? PPO’s critic estimates the return expected from a state. GiGPPO additionally supplies the critic with a visually similar state from another rollout of the same task, that neighbor’s observed return, and the screens’ perceptual-hash distance. The critic learns how to use this evidence rather than simply averaging neighboring returns.

Reading the chart. At update 168, UINavBench success is 16.7% for the SFT starting model, 17.6% for terminal-only PPO, 25.5% for dense-reward PPO, and 30.4% for GiGPPO. GiGPPO is above Gemini 3.5 Flash’s 28.4% and below Gemini 3.1 Pro’s 32.4% in this comparison.

Read paper ↗

PICARD MEETS DIFFUSION

High-accuracy and dimension-free sampling with diffusions

Conceptual illustration of diffusion sampling and Picard refinement within a time window; the dots mark polynomial collocation nodes.

Diffusion sampling turns noise into data by solving a differential equation driven by the score, the gradient of a noisy distribution’s log density. Numerical error matters: earlier sampling guarantees require polynomially many iterations in the inverse target error, 1/ε.

Picard iteration refines a candidate path. On each short time window, it evaluates the equation’s vector field along the current path, integrates those values, and uses the result as the next path estimate. Our solver makes this finite computation possible through polynomial collocation: evaluate at selected times, interpolate with a low-degree polynomial, and integrate that polynomial. The illustration’s inset depicts refinement within one window; the sampler chains multiple windows together.

Why does this give high accuracy? We prove that the score along the reverse process admits an accurate low-degree approximation. Combined with Picard contraction, this yields iteration complexity polylogarithmic in 1/ε.

The guarantee assumes a bounded-support distribution smoothed by Gaussian noise and sufficiently accurate, regular score estimates. “Dimension-free” means dependence through the effective support radius rather than an explicit polynomial in ambient dimension; that radius can still grow with dimension. This is a theoretical sampling guarantee.

Read paper ↗

BLACKWELL-GUIDED LOGIT SEARCH

Blackwell-Guided Logit Search (BLS): Optimizing Multi-Turn Jailbreaks with Subliminal Steering

Conceptual BLS workflow showing vector feedback from logit probes and Blackwell approachability prioritizing unmet objectives.

How can we measure a model’s changing behavior when its replies keep looking like the same refusal? Standard automated red-teaming often generates a full response and asks another model to judge it. That is expensive, and nearly identical refusals can hide differences between candidate conversations.

BLS measures those differences through next-token scores. A set of diagnostic questions produces a vector of signals about the target model’s behavior. These signals can distinguish candidate continuations without generating a full target reply for each one or calling an external judge during selection. The method builds conversations incrementally; the target still replies between completed turns, and a judge evaluates the final outcome.

Why Blackwell approachability? Each diagnostic signal has a desired range. The selection rule tracks the running-average vector and gives more weight to objectives that remain furthest outside their ranges. Objectives already satisfied receive zero weight. This lets the search respond to several aspects of behavior rather than relying on one fixed score.

What do the experiments show? On one hardened 72B medical-role scenario, vector BLS reaches a 37.3% break rate, compared with 7.1% for TAP and 0% for PAIR at matched attacker-token budgets. Across the reported scenario comparisons, it uses 22–37 times fewer target-generated replies; this does not eliminate the cost of logit probes. Separately, the selected scalar BLS variant reaches 61% on JailbreakBench, 74% on HarmBench, and 70% on StrongREJECT. Those benchmark results come from a different variant and include post hoc configuration selection.

The paper also finds that selecting unchanged sentences from a pretraining corpus can induce persona shifts despite little apparent connection between the sentences and the resulting behavior. These are controlled experiments on open-weight models, and BLS requires access to next-token scores. The findings highlight behavior that surface-text judgments alone can miss.

Manuscript · Submitted to ICLR 2027