Research on alignment, model efficiency, and interpretability. My name is bolded where applicable.
Preprints & Under Review
CATPO: Critique-Augmented Tree Policy Optimization
Ankur Dahiya, Umang Goyal, Ayush
Submitted to COLM 2026 · Under review
▸ Abstract
We propose CATPO, a novel reinforcement learning alignment framework that augments tree-search policy optimization with structured critique feedback. Standard RLHF pipelines rely on scalar reward signals that provide limited guidance for complex reasoning tasks. CATPO addresses this by integrating a critique model that generates fine-grained, natural-language evaluations at each node of the policy search tree. This critique signal enables more informed exploration, allowing the policy to prune unpromising branches early and focus computational resources on high-potential trajectories. Through experiments on instruction-following and reasoning benchmarks, we demonstrate that CATPO produces policies that are better aligned with human preferences, achieving improved win-rates against baseline RLHF and DPO approaches while requiring fewer rollout samples. Our analysis reveals that the critique-augmented tree structure enables emergent self-correction behavior, where the policy learns to identify and recover from its own reasoning errors during generation.
We draw a principled analogy between biological natural selection and neural network pruning to develop an evolutionary pruning framework for large language models. Our method treats subnetworks as individual organisms within a population, defining a fitness function based on task-specific performance and parameter efficiency. Through iterative cycles of evaluation, selection, crossover, and mutation, we evolve compact model architectures that emerge from dense pretrained networks. Unlike conventional magnitude or gradient-based pruning methods that operate in a single pass, our evolutionary approach explores a diverse space of possible sparse architectures and discovers non-obvious pruning patterns that static methods miss. Experiments on LLaMA-family models show that our evolved subnetworks retain 92–97% of the dense model's performance at 40–60% sparsity, consistently outperforming structured and unstructured pruning baselines across language modeling, commonsense reasoning, and instruction-following tasks.