Pritish Saha

IIT Kharagpur / AI research / LLM systems

I study when model behavior is grounded in internal computation.

I am a 4th-year Dual Degree student at IIT Kharagpur: B.Tech. (Hons.) in Manufacturing Science and Engineering and M.Tech. in Industrial Engineering and Management. My current work sits around reinforcement learning for reasoning models, representation-level supervision, geometry-aware PEFT, learned memory, and mechanistic latent-state analysis.

The question I keep coming back to is simple: if an agent gets the right answer, did it use a robust internal strategy or just a shortcut that happened to work? I like problems where behavior, representation, and causal interventions all have to tell the same story.

CV Email IIT KGP mail LinkedIn GitHub X Scholar Hugging Face
+4.70ppLaViDA nearest-expert gain over GRPO on MATH-500 n=8, seed 0
R^2=0.9826D Bayesian posterior decoded from transformer residual streams
262Klearned memory tokens scaled with Mixture of Chapters routing
25-80%trainable parameter reduction targeted by GRIT-style PEFT

Research

Current threads

The common theme is internal structure: hidden states as belief states, LoRA updates as geometry, memory as an addressable substrate, and RL as distribution shaping rather than only final-answer reward.

BTP / Complex Networks Research Lab, IIT Kharagpur / Prof. Pawan Goyal

LaViDA: Latent Visitation Distribution Alignment for Mathematical Reasoning

LaViDA asks whether outcome-only GRPO leaves useful reasoning structure on the table. Correct rollouts are projected through a frozen latent encoder, then aligned toward verified expert traces. The current result is deliberately honest: simple nearest-expert alignment is promising under sampling, while the learned chi-square critic is not yet reward-aligned.

  • Built Qwen2.5-Math-7B GRPO training with LoRA-r64, vLLM, FlashAttention, and single-H100 rollout scaling.
  • Constructed self and Oracle-augmented expert pools: 8,963 self traces plus 3,354 filtered 72B traces embedded on the 7B manifold.
  • D_OracleAug ties GRPO on greedy MATH-500 but improves n=8 mean correctness by +4.70pp, p=0.0069.
  • On the harder MATH-500 L4-5 subset, the same branch improves n=8 mean correctness by +5.77pp, p=0.0429.
  • Chi-square density-ratio matching was null; the post-BTP direction is CR-LaViDA, a prompt-exclusive InfoNCE reformulation.

BTP slides

MATH-500 n=8 mean correctness

A_2000 GRPO
74.88%
D_OracleAug
79.57%
SFT_Oracle
79.45%

Seed-0 read: SFT wins greedy; D_OracleAug and SFT are statistically indistinguishable under n=8 sampling.

Cambridge AI Safety Hub / MARS 4.0 / Prof. Fernando Rosas

Bayesian mixed-state geometry in transformers

I study whether small causal transformers trained on hierarchical stochastic generators linearly encode the exact Bayesian belief state, including slow hidden drivers and fast transducer dynamics.

  • Implemented hierarchical HMM/transducer processes with exact 6-state Bayesian filters.
  • Decoded the full 6D posterior from residual streams with R^2=0.982 and MSE=0.0043.
  • Used shuffle, untrained, temporal split, and subspace-overlap controls; driver/transducer separation reached z=-4.69.
MSP fractal compared with learned transformer geometry
Ground-truth mixed-state geometry vs. learned residual-stream geometry.
Layerwise belief geometry decoded from residual stream
Layerwise probes show belief geometry becoming cleaner across depth.
RAAPID INC / Prof. Amitava Das / first-author paper

GRIT: geometry-aware PEFT

GRIT treats adapter updates as a geometric object rather than just a bag of low-rank parameters. It uses rank-space K-FAC, Fisher-guided reprojection, dynamic rank adaptation, and high-rank-to-low-rank compression to make LoRA-style updates more sample-efficient and less drift-prone.

GRIT pipeline from LoRA update through K-FAC natural gradient and neural reprojection
GRIT pipeline: estimate rank-space curvature, precondition the LoRA update, then reproject into the top Fisher eigenspace.
  • Uses a train-high, compress-low setup: train at r=64, then ship a Fisher-compressed adapter at roughly r=12-20.
  • Mid-training compaction is basis change, not freezing: it preserves B@A up to a reconstruction check and only fires under conservative guards.
  • Authored fused Triton kernels for covariance fusion, GPU-side Cholesky inversion, and batched preconditioning.
  • Built a dual-stream CUDA pipeline that overlaps stale-by-one K-FAC inversion/eigendecomposition with the next training step.
  • Targets competitive generative/NLU performance while reducing trainable parameters by 25-80%.
GRIT parameter update geometry compared with LoRA
GRIT concentrates LoRA updates into a tighter curvature-aligned subspace.
GRIT layerwise parameter update ablation
Layerwise update footprint comparison: LoRA, LoRA plus K-FAC, and GRIT.

Paper

ICLR 2026 NFAM Workshop

Mixture of Chapters: learned memory inside transformers

MoC adds a learned internal memory bank that transformer layers query through cross-attention. To scale beyond dense memory access, the bank is partitioned into chapters and a router selects a sparse subset per input sequence. The point is to give models an explicit, trainable memory substrate rather than only implicit parametric storage or external text retrieval.

  • Scales to 262,208 learned memory tokens with 4,097 chapters and sparse top-k routing.
  • Outperforms iso-FLOP vanilla transformer baselines during pretraining.
  • Shows stronger retention under heavy instruction fine-tuning, suggesting memory can reduce interference across phases.

Paper OpenReview Code

Mixture of Chapters architecture diagram
Chapter-routed memory cross-attention.
Mixture of Chapters pretraining loss curve
Pretraining validation loss under iso-FLOP comparison.
Instruction fine-tuning validation loss comparison
Instruction fine-tuning retention comparison.
Medical LLM safety

FROST and antidistillation

I benchmarked unsafe behavior in medical LLM settings and worked on token-level decoding defenses for teacher-side antidistillation. The current direction treats defended generation as a utility-preserving control problem: keep answers useful, but reduce how learnable each trace is for a black-box student.

  • Benchmarked MedGemma on CARES-18K with 86% attack success rate under unsafe prompt settings.
  • Implemented token-level decoding machinery for Fisher-geometric, KL-constrained antidistillation experiments.
Safe Gen-AI course / IIT Kharagpur

Safety fundamentals and assignments

I took CS60216: Safety Fundamentals of Generative AI, covering PEFT/LoRA, RLHF/DPO, jailbreaks, model editing, mechanistic interpretability, scalable oversight, activation patching, watermarking, and training-time alignment.

Course pageAssignments

Publications

Papers

Systems and competitions

Applied work

I keep these on the page because they reflect how I work: build the pipeline, make the evaluation honest, then optimize the bottleneck until the system actually runs.

GenAI analytics dashboard

Runner-up, General Championship Data Analytics, IIT Kharagpur

Captained a full-stack NLQ analytics dashboard for Frammer AI with LangGraph, self-healing SQL, KPI labs, and Gaussian-anchored synthetic star-schema evaluation.

Proof

Amazon ML Challenge 2025

40.8 SMAPE

Stacked Qwen2.5-VL-3B SFT with LightGBM over CLIP/text features; used offline tensorization, WebDataset, 4-bit QLoRA, Pseudo-Huber loss, and monotonic constraints.

Code

American Express Campus Challenge

National Finalist, Decision Science Track

Built a 3-stage GBDT-Transformer ranking ensemble with 3k+ leakage-free temporal features and a listwise Transformer trained on GBDT residuals; final MAP 0.59.

Code

Background

Education

Indian Institute of Technology, Kharagpur
B.Tech. (Hons.) in Manufacturing Science and Engineering and M.Tech. in Industrial Engineering and Management, 2022-2027.

Coursework and self-study include Safety Fundamentals of Generative AI, Operations Research, Probability and Statistics, Linear Algebra, Stanford CS229, Stanford CS230, LLM Agents MOOC, Algozenith, and Summer Analytics.

Research taste. I am drawn to interactive agent learning, goal-conditioned RL, state abstractions, representation learning, and multimodal agents that can transfer useful structure across settings instead of memorizing isolated tasks.

Stack

Tools I use

PythonC/C++CUDATritonPyTorchJAXPenzaiTransformersPEFT/LoRATRLvLLMFlashAttention-2bitsandbytesFastAPIDockerLinuxWebDatasetLangGraphChromaDBGit

I am open to research collaborations around RL for reasoning, mechanistic supervision, latent-state interpretability, efficient adaptation, learned memory, and agents that learn useful abstractions from limited interaction.

References and letters from Prof. Pawan Goyal and Prof. Amitava Das are available privately on request. Outside research: Codeforces Pupil, interhall football and water polo, karate black belt, and NSS volunteering.