LLM safety · alignment · test-time reasoning

Vaibhav Singh

I work on the space between a model's first token and its last — steering decoding away from unsafe completions, choosing which examples a model gets to see, and intervening early in a reasoning trace when it starts to go wrong.

In Tokyo, at SB Intuitions. Previously Axelspace, Yahoo! JAPAN and ShareChat. Engineering Physics, IIT Bombay.

Published at CVPR·ICML·EMNLP
Working as AI/ML Data Scientist
Founded Lombard Standard·Antler

Behind this page: Conway's Game of Life. Click anywhere to perturb it — a few cells changes everything.

Research

CVPR · ICML · EMNLP · 2025–2026

One question runs through all of it: a model is already generating — what is the cheapest thing you can change, and when, to make the output better? Each diagram below is drawn in the same grammar.

the process running the intervention what happens without it
safe jailbroken step 2
ICML 2026

Safety Recovery in Reasoning Models Is Only a Few Early Steering Steps Away

Soumya Suvra Ghosal, Souradip Chakraborty, Vaibhav Singh, Furong Huang, Dinesh Manocha, Amrit Singh Bedi

RL post-training makes multimodal reasoning models better at thinking and worse at refusing. SafeThink watches the reasoning trace with a safety reward model and, only when a threshold is crossed, injects a short corrective prefix. Safety is treated as a constraint to satisfy, not a score to maximise.

Attack success down 30–60% across six open MLRMs and four jailbreak benchmarks.
LlamaV-o1 on JailbreakV-28K: 63.33% → 5.74%, with reasoning accuracy intact.

Abstract

Reinforcement-learning-based post-training for explicit chain-of-thought reasoning improves the reasoning ability of multimodal large reasoning models, but it can simultaneously degrade safety alignment and raise jailbreak success rates. We propose SafeThink, a lightweight inference-time defense that treats safety recovery as a satisficing constraint rather than a maximization objective: it monitors the evolving reasoning trace with a safety reward model and conditionally injects an optimized short corrective prefix only when the safety threshold is violated. Across six open-source MLRMs and four jailbreak benchmarks, SafeThink reduces attack success rates by 30–60% while preserving reasoning performance. Empirically, safety recovery is often only a few steering steps away: intervening in the first one to three reasoning steps typically suffices to redirect the full generation toward safe completions.

BibTeX

@inproceedings{ghosal2026safethink,
  title     = {Safety Recovery in Reasoning Models Is Only a Few Early Steering Steps Away},
  author    = {Ghosal, Soumya Suvra and Chakraborty, Souradip and Singh, Vaibhav and
               Huang, Furong and Manocha, Dinesh and Bedi, Amrit Singh},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning (ICML)},
  year      = {2026}
}
solved-example corpus reasoning trace →
EMNLP 2026 Findings · first author

ThinkRetrieve: Retrieval-Augmented Reasoning Traces for Test-Time Scaling

Vaibhav Singh, Soumya Suvra Ghosal, Sarvesh Gharat, Soumyabrata Pal, Ramasuri Narayanam, Dinesh Manocha

Longer chains of thought stop paying off once errors start compounding. ThinkRetrieve pulls solved examples into the trace at each reasoning step, so retrieval supplies procedure — how to attack the next move — rather than another fact.

Up to 60% relative gain on AIME 2025, consistent across five reasoning models from 1.5B to 8B.

Abstract

Test-time scaling in large reasoning models is limited by diminishing returns: extending the reasoning chain accumulates errors rather than correcting them. We introduce ThinkRetrieve, which augments the reasoning traces of LRMs with dynamically retrieved solved examples at each reasoning step. By incorporating relevant exemplars from a corpus of problem–solution pairs directly into the thinking process, the method supplies procedural guidance rather than purely factual context. Across five reasoning models spanning 1.5B to 8B parameters and multiple benchmarks, ThinkRetrieve gives consistent improvements, including up to 60% relative gains on AIME 2025.

BibTeX

@inproceedings{singh2026thinkretrieve,
  title     = {ThinkRetrieve: Retrieval-Augmented Reasoning Traces for Test-Time Scaling},
  author    = {Singh, Vaibhav and Ghosal, Soumya Suvra and Gharat, Sarvesh and
               Pal, Soumyabrata and Narayanam, Ramasuri and Manocha, Dinesh},
  booktitle = {Findings of the Association for Computational Linguistics: EMNLP 2026},
  year      = {2026}
}
k-NN — near, but nearly identical KITE — spread, submodular, kernelized query
EMNLP 2026 Findings · first author

KITE: Kernelized and Information Theoretic Exemplars for In-Context Learning

Vaibhav Singh, Soumya Suvra Ghosal, Kapu Nirmal Joshua, Soumyabrata Pal, Sayak Ray Chowdhury

The context window is small, so example choice is the whole game. Nearest-neighbour retrieval like KATE returns five near-duplicates. KITE reframes selection as minimising prediction error for this one query, derives an approximately submodular objective, and picks greedily — with a guarantee.

Kernel trick for high-dimensional feature spaces + an optimal-design regulariser that forces diversity.
Beats standard retrieval across a suite of label-scarce classification tasks.

Abstract

In-context learning has emerged as a powerful paradigm for adapting large language models to new and data-scarce tasks using only a few carefully selected task-specific examples presented in the prompt. However, given the limited context size of LLMs, a fundamental question arises: which examples should be selected to maximize performance on a given user query? While nearest-neighbor-based methods like KATE have been widely adopted, they suffer from well-known drawbacks in high-dimensional embedding spaces, including poor generalization and a lack of diversity. We study example selection from a principled, information-theory-driven perspective. We model an LLM as a linear function over input embeddings and frame example selection as a query-specific optimization problem: selecting a subset of exemplars from a larger bank that minimizes the prediction error on a specific query. We derive a principled surrogate objective that is approximately submodular, enabling a greedy algorithm with an approximation guarantee, and further (i) incorporate the kernel trick to operate in high-dimensional feature spaces without explicit mappings, and (ii) introduce an optimal-design-based regularizer to encourage diversity. Empirically we show significant improvements over standard retrieval methods across a suite of classification tasks.

BibTeX

@inproceedings{singh2026kite,
  title     = {KITE: Kernelized and Information Theoretic Exemplars for In-Context Learning},
  author    = {Singh, Vaibhav and Ghosal, Soumya Suvra and Joshua, Kapu Nirmal and
               Pal, Soumyabrata and Chowdhury, Sayak Ray},
  booktitle = {Findings of the Association for Computational Linguistics: EMNLP 2026},
  year      = {2026}
}
safe reward gate controlled decoding
CVPR 2025

IMMUNE: Improving Safety Against Jailbreaks in Multi-modal LLMs via Inference-Time Alignment

Soumya Suvra Ghosal, Souradip Chakraborty, Vaibhav Singh, Tianrui Guan, Mengdi Wang, Alvaro Velasquez, Ahmad Beirami, Furong Huang, Dinesh Manocha, Amrit Singh Bedi

Safety training on the base model does not survive contact with an adversarial image. IMMUNE leaves the weights alone and moves the defense into decoding: a safe reward model scores continuations as they are generated, so unsafe branches never get chosen. With a formal guarantee on the aligned distribution it converges to.

On LLaVA-1.6, attack success down 57.82% vs. the base MLLM and 16.78% vs. the previous best defense — general capability unchanged.

Abstract

With the widespread deployment of multimodal large language models for visual reasoning, improving their safety has become crucial. Recent studies show that despite safety-alignment training, these models remain vulnerable to jailbreak attacks. We introduce IMMUNE, an inference-time defense framework that leverages a safe reward model through controlled decoding to defend against jailbreak attacks. Beyond providing a mathematical characterisation of the approach, extensive evaluation across recent MLLMs and attack strategies demonstrates its effectiveness: on LLaVA-1.6, IMMUNE reduces the attack success rate by 57.82% and 16.78% compared to the base MLLM and the state-of-the-art defense strategy respectively, while preserving the model's original capabilities.

BibTeX

@inproceedings{ghosal2025immune,
  title     = {Immune: Improving Safety Against Jailbreaks in Multi-modal LLMs
               via Inference-Time Alignment},
  author    = {Ghosal, Soumya Suvra and Chakraborty, Souradip and Singh, Vaibhav and
               Guan, Tianrui and Wang, Mengdi and Velasquez, Alvaro and Beirami, Ahmad and
               Huang, Furong and Manocha, Dinesh and Bedi, Amrit Singh},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and
               Pattern Recognition (CVPR)},
  year      = {2025}
}
high-resource bank preferred rejected margin
EMNLP 2025 Findings

RELIC: Enhancing Reward Model Generalization for Low-Resource Indic Languages with Few-Shot Examples

Soumya Suvra Ghosal, Vaibhav Singh, Akash Ghosh, Soumyabrata Pal, Subhadip Baidya, Sriparna Saha, Dinesh Manocha

Open multilingual reward models are trained on high-resource preference data, so their reward signal for Indic languages is close to noise. RELIC trains a retriever with a pairwise ranking objective to pull in-context examples from auxiliary high-resource languages — specifically the ones that sharpen the gap between a preferred and a rejected response.

Bodo with a LLaMA-3.2-3B reward model: +12.81% accuracy over zero-shot and +10.13% over the best prior selection method, across PKU-SafeRLHF, WebGPT and HH-RLHF.

Abstract

Reward models are essential for aligning large language models with human preferences, yet most open-source multilingual reward models are trained primarily on preference data in high-resource languages, producing unreliable reward signals for low-resource Indic languages. We propose RELIC, an in-context learning framework for reward modeling in low-resource Indic languages. RELIC trains a retriever with a pairwise ranking objective to select in-context examples from auxiliary high-resource languages that most effectively highlight the distinction between preferred and less-preferred responses. Experiments on PKU-SafeRLHF, WebGPT and HH-RLHF show that RELIC significantly improves reward model accuracy for low-resource Indic languages; on Bodo, using a LLaMA-3.2-3B reward model, RELIC achieves 12.81% and 10.13% improvements over zero-shot prompting and the state-of-the-art example selection method respectively.

BibTeX

@inproceedings{ghosal2025relic,
  title     = {RELIC: Enhancing Reward Model Generalization for Low-Resource Indic
               Languages with Few-Shot Examples},
  author    = {Ghosal, Soumya Suvra and Singh, Vaibhav and Ghosh, Akash and
               Pal, Soumyabrata and Baidya, Subhadip and Saha, Sriparna and Manocha, Dinesh},
  booktitle = {Findings of the Association for Computational Linguistics: EMNLP 2025},
  year      = {2025}
}

Frequent collaborators: Soumya Suvra Ghosal, Soumyabrata Pal, Souradip Chakraborty, Dinesh Manocha and Amrit Singh Bedi (University of Maryland · Adobe Research · UCF).

Experience

Tokyo, Japan

Research is one half. The other half is shipping models that have to survive real traffic.

PresentTokyo, JP

Senior Data Scientist

SB Intuitions · Creative Vision Team, Multimodal AI Dept.

Image generation models, and world models trained from scratch. Building the datasets they learn from, and the evaluation that tells you whether any of it actually worked.

  • Image Generation
  • World Models
  • Training from Scratch
  • Dataset Curation
  • Model Evaluation
PreviouslyTokyo, JP

Senior Data Scientist

Computer vision for microsatellite earth observation — models that read satellite imagery, and the infrastructure that kept them running.

  • Computer Vision
  • Satellite Imagery
  • PyTorch
  • Kubernetes

Projects

Open source · before the papers

Things built for the sake of building them. Most predate the research, all are on GitHub.

Offscreen

Conferences · Tokyo

Where the work happens, and where it doesn't.

Get in touch

Open to research collaboration

If you are working on inference-time control, reward models, or anything that makes a model safer without retraining it — I would like to hear about it. Papers, code and datasets above are all public; ask me anything about them.

gen 0