Sayak Chakrabarti

সায়ক চক্রবর্ত্তী

PhD Candidate in Computer Science, Columbia University

I'm a PhD candidate in Computer Science at Columbia University, fortunate to be advised by Josh Alman and Toniann Pitassi. I work on getting large language models to learn more from feedback and to run on less compute. My research interests are post-training, long-horizon reinforcement learning, PPO/DPO/GRPO, quantization, and efficient Transformer architectures.

Recently, I interned at Zoom AI, where I worked on credit assignment for critic-free long-horizon RL, which led to RELACE.

Earlier, I did a dual B.Tech–M.Tech at IIT Kanpur under Nitin Saxena, who shaped how I approach research. My thesis there on polynomial root-finding received the Best BT-MT Thesis Award. I've also spent time at MPI-SWS and TIFR.

Publications

Also on Google Scholar

Machine learning

RELACE: Retrospective Likelihood-based Action Credit Estimation for Long-Horizon Language Agents

Sayak Chakrabarti, Sathish Reddy Indurthi

Under review · 2026 Post-training Long-horizon RL

Critic-free, step-level credit for multi-turn agents. Each executed action is scored under its original context and under an outcome-augmented one, and the shift in likelihood reweights returns among actions from equivalent states. With Qwen2.5-1.5B, RELACE reaches 96.4% on ALFWorld and 79.4% on WebShop, ahead of GiGPO by 5.5 and 5.6 points.

Group Relative Policy Optimization (GRPO) avoids a separate critic by estimating advantages from rollout groups. For multi-turn agents, however, trajectory-level supervision provides coarse, noisy credit: terminal rewards do not locate errors and can penalize useful actions alongside mistakes. Group-in-Group Policy Optimization (GiGPO) and subsequent methods refine supervision through state-conditioned comparisons, but their credit estimates remain sensitive to downstream decisions and outcomes. We introduce RELACE (Retrospective Likelihood-based Action Credit Estimation), a critic-free framework that integrates retrospective action assessment with state-conditioned advantage estimation. RELACE evaluates executed actions through teacher-forced likelihood scoring under both their original contexts and outcome-augmented contexts. Comparing these likelihoods yields a trajectory-normalized retrospective factor that captures outcome-dependent changes in action plausibility, rather than hindsight plausibility alone. We use this factor to reweight discounted task returns and construct local advantages by comparing weighted returns among actions from equivalent states within a task. This couples retrospective relevance with observed reward, producing fine-grained credit that complements trajectory-level GRPO supervision. Temporal smoothing and success-protecting masking further stabilize the local signal. RELACE requires neither auxiliary value nor reward models nor additional autoregressive rollouts for credit estimation. Experiments on ALFWorld and WebShop with Qwen2.5-1.5B-Instruct and Qwen2.5-7B-Instruct demonstrate substantial improvements over GRPO, GiGPO, and HCAPO. With the 1.5B model, RELACE achieves 96.35% success on ALFWorld and 79.43% on WebShop, surpassing GiGPO by 5.47 and 5.60 percentage points, respectively.

@article{chakrabarti2026relace,
  title   = {{RELACE}: Retrospective Likelihood-based Action Credit Estimation for Long-Horizon Language Agents},
  author  = {Chakrabarti, Sayak and Indurthi, Sathish Reddy},
  journal = {arXiv preprint arXiv:2610.07349},
  year    = {2026}
}

Every Bit Counts: A Theoretical Study of Precision–Expressivity Tradeoffs in Quantized Transformers

Sayak Chakrabarti, Toniann Pitassi, Josh Alman

Under review · 2026 Quantization

For every precision p, we construct an equality-style function that a one-layer softmax Transformer can compute with p bits but provably cannot compute with p − 1. So tasks built on exact matching are especially sensitive to quantization, and the threshold is sharp. The proofs pair explicit finite-precision constructions with communication-complexity lower bounds, and PTQ/QAT experiments show the predicted accuracy drop.

Quantization reduces the numerical precision of Transformer computations and is widely used to accelerate inference, yet its effect on expressivity remains poorly characterized. We demonstrate a fine-grained theoretical tradeoff between expressivity and precision: for every p we exhibit a function Γ, inspired by the equality function, and prove that a one-layer softmax Transformer can compute Γ with p bits of precision, but not with p − 1 bits of precision.

This result concretely explains the widely observed phenomenon of empirical loss of expressivity when quantization is used. Practically, it suggests that tasks requiring equality-like comparisons (exact match, membership, etc.) are especially sensitive to quantization. Dropping even one bit can cross a threshold where the model cannot represent the needed comparison reliably. Thus, it paves the way for developing heuristics that will help practitioners choose how much quantization is possible: the precision should be chosen as a function of the length of equality to be checked for the specific task.

Our proofs combine explicit finite-precision Transformer constructions with communication-complexity lower bounds, yielding a tight “one-bit” threshold.

@article{chakrabarti2026everybit,
  title   = {Every Bit Counts: A Theoretical Study of Precision-Expressivity Tradeoffs in Quantized Transformers},
  author  = {Chakrabarti, Sayak and Pitassi, Toniann and Alman, Josh},
  journal = {arXiv preprint arXiv:2602.02707},
  year    = {2026}
}

Poly-attention: A General Scheme for Higher-Order Self-Attention

Sayak Chakrabarti, Toniann Pitassi, Josh Alman

ICLR 2026 Efficient architectures

A general family of higher-order attention mechanisms that contains tensor and Strassen attention as special cases, with matching algorithms and lower bounds for exact and approximate computation. Its tree-attention subclass runs in quadratic time yet provably composes any fixed number of functions, which prior mechanisms needed superquadratic time to do even for two.

The self-attention mechanism, at the heart of the Transformer model, is able to effectively model pairwise interactions between tokens. However, numerous recent works have shown that it is unable to perform basic tasks involving detecting triples of correlated tokens, or compositional tasks where multiple input tokens need to be referenced to generate a result. Some higher-dimensional alternatives to self-attention have been proposed to address this, including higher-order attention (Sanford et al., 2023) and Strassen attention (Kozachinskiy et al., 2025), which can perform some of these polyadic tasks in exchange for slower, superquadratic running times.

In this work, we define a vast class of generalizations of self-attention, which we call poly-attention mechanisms. Our mechanisms can incorporate arbitrary higher-order (tensor) computations as well as arbitrary relationship structures between the input tokens, and they include the aforementioned alternatives as special cases. We then systematically study their computational complexity and representational strength, including giving new algorithms and matching complexity-theoretic lower bounds on the time complexity of computing the attention matrix exactly as well as approximately, and tightly determining which polyadic tasks they can each perform. Our results give interesting trade-offs between different desiderata for these mechanisms, including a tight relationship between how expressive a mechanism is, and how large the coefficients in the model may be so that the mechanism can be approximated in almost-linear time.

Notably, we give a new attention mechanism which can be computed exactly in quadratic time, and which can perform function composition for any fixed number of functions. Prior mechanisms, even for just composing two functions, could only be computed in superquadratic time, and our new lower bounds show that faster algorithms for them are not possible.

@inproceedings{chakrabarti2026polyattention,
  title     = {Poly-attention: A General Scheme for Higher-Order Self-Attention},
  author    = {Chakrabarti, Sayak and Pitassi, Toniann and Alman, Josh},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2026}
}

Algorithms & computational algebra

Before moving to ML, I worked on algorithms for polynomials over finite rings.

Solving polynomial systems over non-fields and applications to modular polynomial factoring

(α-β) Sayak Chakrabarti, Ashish Dwivedi, Nitin Saxena

Journal of Symbolic Computation 125, 2024

We study the problem of solving a system of n-variate m polynomials over the ring of integers modulo a prime-power pk. The problem over finite fields is well studied in varied parameter settings. For small characteristic p = 2, Lokshtanov et al. (SODA'17) initiated the study, for degree d = 2 systems, to improve the brute force complexity of O(2n) to O(20.8765n), which currently is improved to O(20.6943n) in Dinur (SODA'21). For large p but constant n, Huang and Wong (FOCS'96) gave a randomized poly(d, m, log p) time algorithm. Note that for growing n, system-solving is known to be intractable even with p = 2 and degree d = 2.

We devise a randomized poly(d, m, log p)-time algorithm to find a root of a given system of m integral polynomials of degrees bounded by d, in n variables, modulo a prime power pk, when n + k is constant. In a way, we extend the efficient algorithm of Huang and Wong (FOCS'96) for system-solving over Galois fields (i.e. characteristic p) to that over Galois rings (i.e. characteristic pk), when k > 1 is constant. The challenge here is to find a lift of singular 𝔽p-roots (exponentially many), as there is no efficient general way known in algebraic geometry for resolving singularities.

Our algorithm has application to factoring univariates over Galois rings. Given f ∈ ℤ[x] and a prime-power pk (k ≥ 2), finding factors of f mod pk has a curious state-of-the-art. It is solved for ‘large’ k by p-adic factoring algorithms (von zur Gathen, Hartlieb, ISSAC'96), but unsolved for ‘small’ k. In particular, no nontrivial factoring method is known for k ≥ 5 (Dwivedi, Mittal, Saxena, ISSAC'19). One issue is that degree-δ factors of f(x) mod pk could be exponentially many, as soon as k ≥ 2. We give the first randomized poly(deg(f), log p)-time algorithm to find a degree-δ factor of f(x) mod pk, when k + δ is constant. Our method has potential application in algebraic coding theory. Algebraic codes, e.g. algebraic-geometric (AG) and Reed–Solomon (RS) codes, are well-studied over Galois rings to achieve better bounds on their parameters.

@article{chakrabarti2024solving,
  title   = {Solving polynomial systems over non-fields and applications to modular polynomial factoring},
  author  = {Chakrabarti, Sayak and Dwivedi, Ashish and Saxena, Nitin},
  journal = {Journal of Symbolic Computation},
  volume  = {125},
  pages   = {102314},
  year    = {2024},
  doi     = {10.1016/j.jsc.2024.102314}
}

An effective description of the roots of bivariates mod pk and the related Igusa's local zeta function

(α-β) Sayak Chakrabarti, Nitin Saxena

International Symposium on Symbolic and Algebraic Computation (ISSAC), 2023

Finding roots of a bivariate polynomial f(x1, x2), over a prime field 𝔽p, is a fundamental question with a long history and several practical algorithms are now known. Effective algorithms for describing the roots modulo pk, k ≥ 2, for any general bivariate polynomial, were unknown until the present paper. The main obstruction is lifting the singular 𝔽p roots to ℤ/pkℤ. Such roots may be numerous and behave unpredictably, i.e., they may or may not lift from ℤ/pjℤ to ℤ/pj+1ℤ.

We give the first algorithm to describe the roots of a bivariate polynomial over ℤ/pkℤ in a practical way. Notably, when the degree of the polynomial is constant, our algorithm runs in deterministic time which is polynomial in p + k. This is a significant improvement over brute force, which would require exploring p2k possible values. Our method also gives the first efficient algorithms for the following problems (which were also open): (1) efficiently representing all the (possibly infinitely many) p-adic roots, and (2) computing the underlying Igusa's local zeta function. We also obtain a new, effective method to prove the rationality of this zeta function.

@inproceedings{chakrabarti2023effective,
  title     = {An effective description of the roots of bivariates mod $p^k$ and the related {Igusa's} local zeta function},
  author    = {Chakrabarti, Sayak and Saxena, Nitin},
  booktitle = {Proceedings of the International Symposium on Symbolic and Algebraic Computation (ISSAC)},
  year      = {2023},
  doi       = {10.1145/3597066.3597115}
}

On Algorithms to Find p-ordering

Aditya Gulati, Sayak Chakrabarti, Rajat Mittal

Conference on Algorithms and Discrete Applied Mathematics (CALDAM), 2021

The concept of p-ordering for a prime p was introduced by Manjul Bhargava (in his PhD thesis) to develop a generalized factorial function over an arbitrary subset of integers. This notion of p-ordering provides a representation of polynomials modulo prime powers, and has been used to prove properties of roots sets modulo prime powers. We focus on the complexity of finding a p-ordering given a prime p, an exponent k and a subset of integers modulo pk.

Our first algorithm gives a p-ordering for a set of size n in time O(nk log p), where the set is considered modulo pk. The subsets modulo pk can be represented succinctly using the notion of representative roots; a natural question would be, can we find a p-ordering more efficiently given this succinct representation. Our second algorithm achieves precisely that: we give a p-ordering in time O(d2k log p + nk log p + nd), where d is the size of the succinct representation and n is the required length of the p-ordering. Another contribution that we make is to compute the structure of roots sets for prime powers pk, when k is small. We explicitly describe all the root sets for p2, p3 and p4.

@inproceedings{gulati2021pordering,
  title     = {On Algorithms to Find p-ordering},
  author    = {Gulati, Aditya and Chakrabarti, Sayak and Mittal, Rajat},
  booktitle = {Conference on Algorithms and Discrete Applied Mathematics (CALDAM)},
  year      = {2021},
  doi       = {10.1007/978-3-030-67899-9_27}
}

Graphon estimation from partially observed network data

Soumendu Sundar Mukherjee, Sayak Chakrabarti

Preprint, 2019

We consider estimating the edge-probability matrix of a network generated from a graphon model when the full network is not observed—only some overlapping subgraphs are. We extend the neighbourhood smoothing (NBS) algorithm of Zhang et al. (2017) to this missing-data set-up and show experimentally that, for a wide range of graphons, the extended NBS algorithm achieves significantly smaller error rates than standard graphon estimation algorithms such as vanilla neighbourhood smoothing (NBS), universal singular value thresholding (USVT), blockmodel approximation, matrix completion, etc. We also show that the extended NBS algorithm is much more robust to missing data.

@article{mukherjee2019graphon,
  title   = {Graphon estimation from partially observed network data},
  author  = {Mukherjee, Soumendu Sundar and Chakrabarti, Sayak},
  journal = {arXiv preprint arXiv:1906.00494},
  year    = {2019}
}

(α-β) denotes alphabetical author order, the convention in theoretical computer science.

Theses & reports My M.Tech thesis and undergraduate research reports are archived on a separate page.

Experience

Research

Education

Teaching & Service

Teaching

  • COMS 3261: Computer Science TheoryTA, Columbia · Spring 2025, Fall 2023
  • ESC101: Fundamentals of ComputingTutor, IIT Kanpur · 2021–22
  • CS203: Probability for Computer ScienceTA, IIT Kanpur · Spring 2021
  • CS202: Logic for Computer ScienceTA, IIT Kanpur · Spring 2021

Reviewing

  • ConferencesICLR 2026, ITCS 2026, SODA 2026, ESA 2024, SOSA 2024, MFCS 2023
  • JournalsInformation and Computation, Journal of Number Theory