Arnav Raj

Reading Multi-Token Answers from a Single Hidden State

Solo-authored. Patch a late-layer hidden state into an early layer of a copy prompt and the model spells the answer out itself. That roughly doubles exact-match accuracy over patching at the layer the state came from. Which is the honest comparison to make. The reason to bother at all: a token lens reads a hidden state one token at a time, so an answer like "North Korea" never comes out whole.

PEBS: Per-rater Empirical-Bayes Shrinkage for RLHF Reward-Model Calibration

Thousands of raters disagree with each other, and the standard recipe averages them into a single number. The average rater is nobody in particular. PEBS calibrates to each individual rater instead: a thin statistical layer that sits on top. Alignment to genuine human preferences improves on the standard pluralistic-alignment benchmarks. The reward model underneath is left exactly as it was, and that part is not decoration. If a calibration method costs you a retraining run, nobody runs it.

Retroactive Advantage Correction: Closed-Form V-Trace Bias Correction for Delay-Aware RLHF

In a production RLHF pipeline the reward tends to land late, held up behind a slow code verifier or a large judge ensemble; sometimes it is just a queue of human review. The standard answers are to wait for it or to drop the sample. Dropping is the wrong call. The feedback does exist, it is only late. What the paper gives you is a correction primitive that keeps training on the delayed signal, closed form, with a proof that it stays unbiased. It improves substantially on the standard wait-or-drop strategies.

KG-MuLQA: Multi-hop Question Answering over Knowledge Graphs for Long-Context Evaluation

A multi-hop benchmark for long-context models, with the questions grown out of knowledge graphs over real financial documents. Financial filings make a boring corpus. They are also long and they are real, and for this benchmark that counts for more. My part was the pipeline: question generation, answer synthesis, and the harness that scores frontier LLMs.

Hyperbolic Geometry of Reasoning: Probing LLM Hidden States

Solo-authored interpretability work on how chain-of-thought reasoning models hold hierarchical structure internally. Flat space is the orthodox assumption here rather than an argued one (mostly, I suspect, because Euclidean distance is what every library hands you for free). Hyperbolic probes recover the hierarchy across the whole network. The Euclidean ones come apart in the deepest layers. Those are the layers that produce the final answer, so reasoning models look to be compressing hierarchy into exactly those representations.

An Eval Harness for LLM-Written Code

Docker-sandboxed harness that runs LLM-written code against hidden tests and hands the failures back for another try. First pass clears a bit under three quarters of the task set (73%).

An RL Agent That Rewrites Python for Speed

PPO agent that edits Python for performance one AST transformation at a time, taking maybe a third off runtime (35%) and a fifth off peak memory (20%) on problems it never trained on.

Hangman AI: transformer plus a dictionary filter

A character-level transformer and hand-written candidate filtering, together solving more than 60% of games.

Personality traits from a user-product graph

A bipartite graph of users and the products they interacted with, with GraphSAGE and GAT variants on top, predicting personality categories and beating the strongest MLP baseline.

Data Search & Retrieval Engine

Lexical search and dense search fail on technical documents in different places, so this one runs an inverted index and a vector index side by side and fuses the two rankings.

Spelling correction that reads the rest of the sentence

A noisy-channel corrector for the errors edit distance can't see, real words sitting in the wrong slot: right about 88% of the time on a curated confusion set.

TotalRecall: Cognitive Health App

A Flutter app for early Alzheimer's, built around memory prompts, mood logging and giving a caregiver a window into both, with Firebase underneath.

Advanced Analytic Tool

The dull half of tabular ML, packaged. Pulls CSV, JSON and SQL into one shape, keeps the pipeline steps swappable, and makes two runs of the same experiment actually comparable.

An MCTS Player for Havannah

Monte Carlo Tree Search with UCB exploration for Havannah, winning over 80% of randomized openings against a tuned RAVE-only baseline.

An SDN Controller That Installs Flows Before the Packets Arrive

A Ryu controller that pushes flow rules down before the traffic arrives, so the throughput number stops measuring the controller.

Reliable Transport in Userspace, Reno vs CUBIC

Reliable transport over UDP, in userspace, with Reno and CUBIC on top of it, to see how reliability and congestion control behave end to end when nothing is hidden in the kernel.

Swapping Pages in xv6

Page swapping for the xv6 teaching kernel. The interesting part is the trap handler, which has to tell a garbage address from a page the kernel moved to disk and still owes you.

I keep coming back to the reward model. It gets trained on thousands of annotators who do not agree, and what everybody does about that is take the mean and start training. I think that is a mistake, and most of what I have written so far is one or another attempt to show it. Final year, CS at IIT Delhi, on the dual degree, bachelor's plus master's (B.Tech + M.Tech).

Through 2026 I did AI evaluation and SFT work for Google DeepMind, as an IC via Topcoder. That meant RSI, Gemini evaluations, and the Colab Bench and ML RL environments. Before that I was at Abundant AI (YC F24), building GPU-enabled RL environments for long-running agents, SWE and data-science ones as well, and the sandbox architecture underneath all of it. Some adversarial ML benchmarks I wrote there ended up in the training pipelines of three of the world's top five AI labs, which I still find slightly absurd.

The newest one, under review at IASEAI, is about reading multi-token answers out of a single hidden state. Token lenses give you one token at a time and cannot spell "North Korea". So instead I patch a late-layer state into an early layer of a copy prompt and let the model say the answer itself. PEBS, at an ICML 2026 workshop, calibrates a reward model to each individual rater. Retroactive Advantage Correction, at the ICML 2026 RLxF workshop, is a closed-form fix for the case where the reward arrives after the step it belongs to, so you can train on it rather than drop it (dropping is the usual move, and it has always looked expensive to me). A third paper probes reasoning models with hyperbolic geometry and finds hierarchy in the deepest layers, where Euclidean probes lose it; that one is at the ICLR 2026 GRaM workshop. I'm also a coauthor on KG-MuLQA, a long-context multi-hop QA benchmark at ACL 2026.

Earlier I was a research intern at Georgia Tech's Financial Services Innovation Lab, and at Harvard's Edge Computing Lab. I co-founded the AI Safety Club at IIT Delhi, and went through BlueDot Impact's alignment course and the ARENA curriculum. Dates and the rest of it are in the resume.

Download Resume (PDF) →