arXiv subjects

Rohit Jain

Publications and source records attributed to Rohit Jain.

At least 19 recordsLinked to original sources

arXiv2026

Composing What Each Teacher Learned: Delta Distillation for Multi-Teacher On-Policy Learning

Multi-teacher on-policy distillation (MOPD) is used in two settings. In common-domain composition, several teachers score each student rollout from one prompt domain and their signals form a single target; in routed-domain distillation, prompts from different domains are assigned to the corresponding specialist. Both settings usually transfer each teacher's endpoint policy, which mixes what post-training changed with preferences inherited from the teacher's base. We introduce $Δ$-MOPD, which transfers each teacher's teacher-minus-base logit shift re-anchored at the student's frozen initialization, and compare it with endpoint supervision in both settings while holding teacher selection fixed. We first expose the mechanism that impedes endpoint transfer: inherited base pull can exceed the post-training shift. Removing it reduces the teacher-term norm ratio and target--student KL. Across our experiments, the results suggest that shift targets are particularly useful when teacher signals are combined at a state. With three composed teachers, $Δ$-MOPD exceeds endpoint composition by $4.11$ Math and $1.95$ five-benchmark points; with two, it matches endpoint accuracy. Under phased routing, it achieves higher mean performance in both phase orders and reduces the observed order gap from $10.50$ to $6.42$ points. Under interleaved routing, where each update involves one teacher, the two targets perform comparably. The phased results provide supporting evidence that the benefit may extend to signals accumulated across training phases. Target construction is thus an independent design axis in MOPD, complementary to teacher selection. The experiment repo could be find at https://github.com/HJSang/Delta-MOPD

cs.LG↗

arXiv2026

Building AI Research Teams for Autonomous Discovery of Linear Attention Architectures

Autonomous agents can increasingly propose ideas, modify code, and run experiments, creating an opportunity to accelerate scientific discovery using teams of collaborative AI researchers. To escape local minima in the discovery, we propose to encourage orthogonal exploration by constructing the team using agents with varying expertise. However, realizing this design requires more than diverse backgrounds. When experiments are costly, a team can test only a limited number of ideas, thereby creating the need to balance exploration and exploitation when selecting a subset of researchers to propose and execute new experiments. Furthermore, each result should shape what the team tries next. We therefore ask how AI researchers with different scientific backgrounds can build on one another's findings and sustain progress under a limited experimental budget. To address these challenges, we build a team of persistent AI researchers, each grounded in a different discipline, who keep private research histories, share experimental evidence, and receive research opportunities through adaptive allocation. We demonstrate the effectiveness of this design by deploying the team to search for linear attention architectures, where it outperforms GDN2 at 1.3B scale on language modeling, common-sense reasoning, and synthetic and real-world retrieval tasks. In controlled comparisons, removing any single organizational component from our design yields markedly weaker discoveries: without disciplinary backgrounds, the search plateaus early; without private histories, researchers make less progress despite identical shared evidence; and uniform selection falls short of adaptive allocation within the same budget. These results suggest that sustained autonomous discovery depends not only on stronger agents, but also on who explores, what they remember, and how research effort is allocated.

cs.CL↗

arXiv2026

Budgeted KV-Cache Repair for Evolving Contexts

Language-model agents revise their working context as information becomes obsolete, but deleting text does not erase its influence on cached computation. Later key-value (KV) states remain conditioned on the deleted content, while recomputing the entire retained suffix can be costly. We introduce budgeted KV-cache repair: recovering task performance by recomputing a user-chosen fraction of the retained suffix. Comparing repair allocations under identical relevance scores and token budgets, we find that refreshing coherent regions can substantially outperform refreshing isolated tokens. This finding motivates Semantic-Unit Refresh (SUR), which allocates the budget to query-relevant text regions using inexpensive embedding scores, without training or an auxiliary transformer forward. To execute these repairs efficiently, a layer-packed executor processes selected states together at each layer while preserving full causal cache access. Across five models, SUR achieves strong recovery in controlled deletion tasks while recomputing only 2-5% of the suffix. On four natural multi-hop QA datasets, it achieves near-full-recompute exact match at a 5% budget with lower post-edit latency. The executor accelerates fragmented refresh by up to 16.1x over interval-wise execution.

cs.AI↗

arXiv2026

SHARPO: Segment-Level Credit Assignment for Agentic Reinforcement Learning

Agentic reinforcement learning (RL) trains a large language model (LLM) to act over long, multi-step interactions. However, a single localized error can cause task failure, while trajectory-level rewards provide limited guidance for assigning credit to individual decisions. To address this limitation, we introduce Segment-level Hindsight Advantage Reweighting for Policy Optimization (SHARPO), a credit-assignment mechanism that refines Group Relative Policy Optimization (GRPO) at the level of environment-facing segments. Inspired by the existing on-policy self-distillation (OPSD) method, SHARPO computes teacher-student log-probability gaps within each segment and uses the resulting signal to compute a bounded multiplier on the GRPO advantage. This multiplier is shared by all tokens within the segment, allowing credit to vary across different segments. With Qwen2.5-7B-Instruct, SHARPO outperforms existing baselines on the ALFWorld and WebShop benchmarks, including GRPO, SDAR, RLSD, and StepOPSD.

cs.LG↗

arXiv2026

Smaller Models, Better Rejects: Preference Distillation Scaling

Preference distillation typically treats a teacher response as preferred and the student's own response as rejected. This assumes that self-generated failures are the most informative negatives and that rejects must come from a model at least as large as the student, making generation costly at scale. We find neither assumption holds: across students from 7B to 72B, smaller frozen models generate rejects with less inference compute yet train stronger students than self-generated rejects, before and after sequence-level knowledge distillation, on code generation and mathematical reasoning. To explain this result, we derive a finite-horizon utility bound for Direct Preference Optimization in a linearized feature model. The bound characterizes favorable reject distributions and motivates three interventions. First, mixing rejects from smaller and student-scale models improves performance as the smaller model's share increases. Second, reassigning rejects to other prompts and shuffling their code tokens still outperform length-matched gibberish, showing that task structure contributes to reject utility. Third, selecting candidates with lower likelihood under the reference policy improves net transfer when higher-likelihood candidates provide less useful contrast. Lower-likelihood selections outperform higher-likelihood ones for every source. These results suggest that effective rejects preserve task structure while limiting coupling to the reference policy, and that smaller frozen models can provide them at low cost.

cs.LG↗

arXiv2026

Practical Online KV Cache Compaction for LLM Agents: An Empirical Study

LLM agents accumulate long trajectories of reasoning steps, tool calls, and environment feedback, making the KV cache a major inference bottleneck. KV cache compaction can reduce this cost, but most prior methods assume a static context where future queries are known or can be approximated offline. Agents instead require online compaction: new information must be compressed before future relevance is known, using proxy queries cheap enough for the inference path. We study online compaction across token eviction (TE) and attention matching (AM), adapting both to compact agent turns and comparing cheap proxy sources such as boundary, repeat-prefill, and delayed future-generation queries. Experiments on BrowseComp-Plus and WideSearch show that immediate compaction often hurts performance, whereas delaying compaction to use the agent's future queries recovers much of the gap. Moreover, TE is often more robust than AM under imperfect proxies. Across models at different scales, TE preserves most of the accuracy while reducing KV cache by 80%, and can improve throughput over the no compaction baseline. These results position proxy-query selection as a core design choice for practical online KV compaction.

cs.CL↗

arXiv2026

Speculate While You Reason: Teaching Agents to Predict Their Next Tool Call via Joint Agent-Speculator RL

Large language model agents often spend substantial wall-clock time waiting for tool call results. Tool-call speculation can hide this latency by predicting and pre-executing an agent's next tool call if the prediction matches the agent's eventual tool call, but existing speculators are typically separate draft models or cached traces that are poorly aligned with the deployed agent's own behavior. We identify this speculator-agent gap and show that the target agent itself is a strong next-call speculator. This points to a simpler design: unifying the agent and speculator within the same model. In this paper, we introduce the self-speculating agent, a single model that both solves tasks in agent mode and predicts its next tool call from partial trajectories in speculator mode, fully reusing prefix KV cache. To enable this dual-mode agent without degrading performance, we propose a joint agent-speculator reinforcement learning method, which derives speculation targets from the agent's own rollouts and alternates agent and speculator updates. Across agentic search QA and conversational tool-use agentic tasks, our method improves average next tool-call Hit@1 from 44.1 to 61.2 for Qwen3-4B and from 48.9 to 66.3 for Qwen3.5-4B, while preserving agent task success.

cs.AI↗

arXiv2025

Agentic Reinforcement Learning for Real-World Code Repair

We tackle the challenge of training reliable code-fixing agents in real repositories, where complex builds and shifting dependencies make evaluation unstable. We developed a verifiable pipeline with success defined as post-fix build validation and improved reproducibility across ~1K real issues by pinning dependencies and disabling automatic upgrades. Building on this, we introduced a scalable simplified pipeline for large-scale reinforcement learning (RL). Using this setup, we supervised fine-tuned Qwen3-32B in the full pipeline and applied RL on top of the SFT model in the simplified environment. The SFT model distilled from GPT-4.1 trajectories performs on par while being 56x smaller, and RL added 7-20% absolute gains under matched train-test conditions. "Thinking mode" was on par or worse in our experiments. Both SFT and RL models failed to generalize across environments, highlighting the importance of matching train-test environments for building reliable real-world code-fixing agents.

cs.LG↗

arXiv2025

Planner-R1: Reward Shaping Enables Efficient Agentic RL with Smaller LLMs

We investigated Agentic RL with large language models on the \textsc{TravelPlanner} benchmark. Our approach, \textsc{Planner-R1}, achieved a \textbf{56.9\%} final-pass rate with only 180 training queries, a $2.7\times$ improvement over GPT-5's $21.2\%$ baseline and the strongest agentic result on the public leaderboard. A central finding was that smaller models (8B) were highly responsive to reward shaping: with dense process-level signals, they reached competitive performance while being $3.5\times$ more compute-efficient and $1.5\times$ more memory-efficient than 32B models. Larger models were more robust under sparse rewards but exhibited smaller relative gains from shaping and higher variance across runs. While curriculum learning offered no significant benefit, shaped rewards consistently amplified learning dynamics, making 8B models the most efficient setting for agentic RL. Crucially, these gains did not come at the cost of overfitting: fine-tuned models mostly maintained or exceeded baseline performance on out-of-domain tasks, including \textsc{Multi-IF}, \textsc{NaturalPlan}, and $τ$-\textsc{Bench}. These results establish reward shaping as a decisive lever for scaling agentic RL, highlight the competitive strength of smaller models, and demonstrate that efficiency can be achieved without sacrificing generalization.

cs.AI↗

arXiv2023

SOTASTREAM: A Streaming Approach to Machine Translation Training

Many machine translation toolkits make use of a data preparation step wherein raw data is transformed into a tensor format that can be used directly by the trainer. This preparation step is increasingly at odds with modern research and development practices because this process produces a static, unchangeable version of the training data, making common training-time needs difficult (e.g., subword sampling), time-consuming (preprocessing with large data can take days), expensive (e.g., disk space), and cumbersome (managing experiment combinatorics). We propose an alternative approach that separates the generation of data from the consumption of that data. In this approach, there is no separate pre-processing step; data generation produces an infinite stream of permutations of the raw training data, which the trainer tensorizes and batches as it is consumed. Additionally, this data stream can be manipulated by a set of user-definable operators that provide on-the-fly modifications, such as data normalization, augmentation or filtering. We release an open-source toolkit, SOTASTREAM, that implements this approach: https://github.com/marian-nmt/sotastream. We show that it cuts training time, adds flexibility, reduces experiment management complexity, and reduces disk space, all without affecting the accuracy of the trained models.

cs.CL↗

arXiv2021

Micro-rheology of a particle in a nonlinear bath: Stochastic Prandtl-Tomlinson model

The motion of Brownian particles in nonlinear baths, such as, e.g., viscoelastic fluids, is of great interest. We theoretically study a simple model for such bath, where two particles are coupled via a sinusoidal potential. This model, which is an extension of the famous Prandtl Tomlinson model, has been found to reproduce some aspects of recent experiments, such as shear-thinning and position oscillations [J. Chem. Phys. {\bf 154}, 184904 (2021)]. Analyzing this model in detail, we show that the predicted behavior of position oscillations agrees qualitatively with experimentally observed trends; (i) oscillations appear only in a certain regime of velocity and trap stiffness of the confining potential, and (ii), the amplitude and frequency of oscillations increase with driving velocity, the latter in a linear fashion. Increasing the potential barrier height of the model yields a rupture transition as a function of driving velocity, where the system abruptly changes from a mildly driven state to a strongly driven state. The frequency of oscillations scales as $(v_0-v_0^*)^{1/2}$ near the rupture velocity $v_0^*$, found for infinite trap stiffness. Investigating the (micro-)viscosity for different parameter ranges, we note that position oscillations leave their signature by an additional (mild) plateau in the flow curves, suggesting that oscillations influence the micro-viscosity. For a time-modulated driving, the mean friction force of the driven particle shows a pronounced resonance behavior, i.e, it changes strongly as a function of driving frequency. The model has two known limits: For infinite trap stiffness, it can be mapped to diffusion in a tilted periodic potential. For infinite bath friction, the original Prandtl Tomlinson model is recovered. We find that the flow curve of the model (roughly) crosses over between these two limiting cases.

cond-mat.soft↗

arXiv2021

Two step micro-rheological behavior in a viscoelastic fluid

We perform micro-rheological experiments with a colloidal bead driven through a viscoelastic worm-like micellar fluid and observe two distinctive shear thinning regimes, each of them displaying a Newtonian-like plateau. The shear thinning behavior at larger velocities is in qualitative agreement with macroscopic rheological experiments. The second process, observed at Weissenberg numbers as small as a few percent, appears to have no analog in macro rheological findings. A simple model introduced earlier captures the observed behavior, and implies that the two shear thinning processes correspond to two different length scales in the fluid. This model also reproduces oscillations which have been observed in this system previously. While the system under macro-shear seems to be near equilibrium for shear rates in the regime of the intermediate Newtonian-like plateau, the one under micro-shear is thus still far from it. The analysis suggests the existence of a length scale of a few micrometres, the nature of which remains elusive.

cond-mat.soft↗

arXiv2020

Regularity results for a penalized boundary obstacle problem

In this paper we are concerned with a two-penalty boundary obstacle problem of interest in thermics, fluid dynamics and electricity. Specifically, we prove existence, uniqueness and optimal regularity of the solutions, and we establish structural properties of the free boundary.

math.AP↗

arXiv2019

Lipper: Synthesizing Thy Speech using Multi-View Lipreading

Lipreading has a lot of potential applications such as in the domain of surveillance and video conferencing. Despite this, most of the work in building lipreading systems has been limited to classifying silent videos into classes representing text phrases. However, there are multiple problems associated with making lipreading a text-based classification task like its dependence on a particular language and vocabulary mapping. Thus, in this paper we propose a multi-view lipreading to audio system, namely Lipper, which models it as a regression task. The model takes silent videos as input and produces speech as the output. With multi-view silent videos, we observe an improvement over single-view speech reconstruction results. We show this by presenting an exhaustive set of experiments for speaker-dependent, out-of-vocabulary and speaker-independent settings. Further, we compare the delay values of Lipper with other speechreading systems in order to show the real-time nature of audio produced. We also perform a user study for the audios produced in order to understand the level of comprehensibility of audios produced using Lipper.

eess.AS↗

arXiv2017

Supersymmetric Schrödinger Operators with Applications to Morse Theory

In 1981 Edward Witten proved a remarkable result where he derived the classical Morse Inequalities using ideas from Supersymmetric (SUSY) Quantum Mechanics. In this regard, one has an example where a Physical Theory has something to say about the underlying Mathematical Structure. The objective of this essay is to understand this classical result from the perspective of Schrödinger Operators. The essay will be divided in four parts. The first part will revisit the classical theory of Morse and recall some of its fundamental results. In the second part, we consider the underlying physical motivations by considering Quantum Mechanics and 0-Dimensional SUSY. The third part will focus on Schrödinger Operators and highlight some of their basic properties. Finally in the last section we will put everything together and present Witten's proof of the Morse Inequalities. Even here we must be completely honest and say that we only follow Witten in proving the Weak Morse Inequalities. The Strong Morse Inequalities are derived using related ideas from Supersymmetry, but mention is made of the techniques used by Witten to get at the Strong Morse Inequalities.

math.DG↗

arXiv2017

The Calabi Conjecture

In this essay we aim to explore the Geometric aspects of the Calabi Conjecture and highlight the techniques of nonlinear Elliptic PDE theory used by S.T. Yau [SY] in obtaining a solution to the problem. Yau proves the existence of a Geometric structure using differential equations, giving importance to the idea that deep insights into geometry can be obtained by studying solutions of such equations. Yau's proof of the existence of a specific class of metrics have found a natural interpretation in recent developments in Theoretical Physics most notably in the formulation of String Theory. We will also attempt to explore the importance of a special case of Yau's solution known as Calabi-Yau Manifolds in the context of holonomy.

math.DG↗

arXiv2017

Small Boxes Big Data: A Deep Learning Approach to Optimize Variable Sized Bin Packing

Bin Packing problems have been widely studied because of their broad applications in different domains. Known as a set of NP-hard problems, they have different vari- ations and many heuristics have been proposed for obtaining approximate solutions. Specifically, for the 1D variable sized bin packing problem, the two key sets of optimization heuristics are the bin assignment and the bin allocation. Usually the performance of a single static optimization heuristic can not beat that of a dynamic one which is tailored for each bin packing instance. Building such an adaptive system requires modeling the relationship between bin features and packing perform profiles. The primary drawbacks of traditional AI machine learnings for this task are the natural limitations of feature engineering, such as the curse of dimensionality and feature selection quality. We introduce a deep learning approach to overcome the drawbacks by applying a large training data set, auto feature selection and fast, accurate labeling. We show in this paper how to build such a system by both theoretical formulation and engineering practices. Our prediction system achieves up to 89% training accuracy and 72% validation accuracy to select the best heuristic that can generate a better quality bin packing solution.

cs.LG↗

arXiv2017

Uniform estimates for the Penalized Boundary Obstacle Problem

In this paper, motivated by a problem arising in random homogenization theory, we initiate the study of uniform estimates for the fractional penalized obstacle problem, $ Δ^{s}u^ε = β_ε (u^ε)$. In particular we consider the penalized boundary obstacle problem, $s = \frac{1}{2}$, and obtain sharp estimates for the solution independent of the penalizing parameter $ε$. This is a generalization of a result due to H. Brezis and D. Kinderlehrer.

math.AP↗