Rahul G. Krishnan
Author directoryAlso published as: Rahul G Krishnan
2026
MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Scaling of Diffusion Language Models
Chen-Hao Chao | Wei-Fang Sun | Junwei Quan | Chun-Yi Lee | Rahul G Krishnan
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Chen-Hao Chao | Wei-Fang Sun | Junwei Quan | Chun-Yi Lee | Rahul G Krishnan
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Masked diffusion models (MDM) exhibit superior generalization when learned using a Partial masking scheme (Prime). This approach converts tokens into sub-tokens and models the diffusion process at the sub-token level. We identify two limitations of the MDM-Prime framework. First, we find that the functional form of the subtokenizer significantly increases the log-likelihood value in the objective when paired with commonly used Byte-Pair-Encoding (BPE) tokenizers. Second, we lack tools to guide the hyperparameter choice of the token granularity in the subtokenizer. To address these limitations, we analyze the optimal design of the subtokenizer that minimizes MDM-Prime training objective and develop MDM-Prime-v2, a masked diffusion language model which incorporates Binary Encoding and Index Shuffling. Our analysis characterizes how token granularity and sub-token entropy influence the training objective and downstream performance, providing principled criteria for subtokenizer design. When extending the model size to 1.1B parameters, MDM-Prime-v2 demonstrates superior average zero-shot accuracy across eight commonsense reasoning benchmarks, outperforming similar-sized baselines including GPT-Neo, OPT, Pythia, Bloom, SMDM, and TinyLLaMA.
Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models
Seo Hyun Kim | Sunwoo Hong | Younwoo Choi | Chen-Hao Chao | Se-Young Yun | Rahul G Krishnan
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Seo Hyun Kim | Sunwoo Hong | Younwoo Choi | Chen-Hao Chao | Se-Young Yun | Rahul G Krishnan
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Masked diffusion language models (dLMs) offer a promising parallel alternative to autoregressive models for complex reasoning. However, they face a unique credit-assignment challenge: reasoning failures typically stem from a few critical, early commitments that cascade downstream. Existing post-training overlooks this dynamic: supervised fine-tuning (SFT) weights all tokens uniformly, while online reinforcement learning (RL) attributes a single trajectory-level reward to every position. We introduce Pivot-SD, an efficient offline framework supervising only these high-impact bottlenecks (pivots). Pivot-SD isolates pivots using an information-gain metric measuring uncertainty reduction over the remaining masked block. It applies bidirectional supervision: successful pivots are reinforced with cross-entropy, while critical errors in failed trajectories are suppressed via targeted unlikelihood, preserving valid surrounding sub-steps. Using only 200 questions and four rollouts each, Pivot-SD improves LLaDA-8B-Instruct over full-sequence SFT and budget-matched diffusion RL baselines across math and code benchmarks.
2025
Agentic Medical Knowledge Graphs Enhance Medical Question Answering: Bridging the Gap Between LLMs and Evolving Medical Knowledge
Mohammad R. Rezaei | Reza Saadati Fard | Jayson L. Parker | Rahul G. Krishnan | Milad Lankarany
Findings of the Association for Computational Linguistics: EMNLP 2025
Mohammad R. Rezaei | Reza Saadati Fard | Jayson L. Parker | Rahul G. Krishnan | Milad Lankarany
Findings of the Association for Computational Linguistics: EMNLP 2025
Large Language Models (LLMs) have greatly advanced medical Question Answering (QA) by leveraging vast clinical data and medical literature. However, the rapid evolution of medical knowledge and the labor-intensive process of manually updating domain-specific resources can undermine the reliability of these systems. We address this challenge with Agentic Medical Graph-RAG (AMG-RAG), a comprehensive framework that automates the construction and continuous updating of Medical Knowledge Graph (MKG), integrates reasoning, and retrieves current external evidence from the MKG for medical QA.Evaluations on the MEDQA and MEDMCQA benchmarks demonstrate the effectiveness of AMG-RAG, achieving an F1 score of 74.1% on MEDQA and an accuracy of 66.34% on MEDMCQA—surpassing both comparable models and those 10 to 100 times larger. By dynamically linking new findings and complex medical concepts, AMG-RAG not only boosts accuracy but also enhances interpretability for medical queries, which has a critical impact on delivering up-to-date, trustworthy medical insights.