Searcharxiv⌕ Search

arXiv subjects

Vahab Mirrokni

Publications and source records attributed to Vahab Mirrokni.

At least 19 recordsLinked to original sources
arXiv2026

A Proof of the Most Informative Boolean Function Conjecture

Let $X$ be uniform on $\{-1,1\}^n$, let $Y$ be obtained by passing its coordinates independently through a binary symmetric channel with crossover probability $p$, and let $g:\{-1,1\}^n\to\{0,1\}$ be a Boolean function. We give a computer-assisted proof of the Courtade--Kumar conjecture $I(g(X);Y)\le1-H_2(p)$, where $H_2$ is binary entropy, with equality attained by dictator functions. The present work builds on the differential-equation method, itself a limiting form of the auxiliary-receiver approach in network information theory using a continuum of degraded receivers. The proof proceeds from a local inequality to a dimension-independent bound on entropy production. Differentiation along the Boolean noise semigroup expresses entropy production as an average of edge costs. The key estimate is therefore an unrestricted Bellman inequality with two mean constraints and two entropy constraints, allowing arbitrary couplings of the edge variables. This paper and its supplement provide the proofs and computational verification records. The document is lengthy because it is designed to be entirely self-contained, deriving all proofs from first principles and reproducing the proofs of cited results. The supplementary material supporting the computer-assisted parts of the proof are available online.

cs.DS↗
arXiv2026

DLB: Distributed Load Balancing at Scale for Generative AI Inference

The reliance on scarce and expensive accelerators such as GPUs and TPUs in modern datacenters places unprecedented demands on backend infrastructure. For workloads characterized by heterogeneous service times and complex multi-stage processing, such as Generative AI, conventional load balancing techniques are often inadequate, relying heavily on costly overprovisioning to maintain service level objectives. This paper introduces DLB, the Distributed Load Balancer, a novel system designed to minimize end-to-end user latency for large-scale, heterogeneous workloads. DLB employs a scalable, distributed design with peer-to-peer probing to maintain real-time visibility into server capacity across large-scale, geographically distributed infrastructure. The system continuously learns latency models to estimate the latency impact of routing decisions, allowing it to effectively manage heterogeneous hardware and diverse model architectures. We provide a novel theoretical analysis of our routing algorithms that establishes their stability and global performance guarantees over time. We also evaluate DLB through extensive simulations, which show substantial gains compared to state-of-the-art load balancing algorithms. Finally, following a 22-month deployment of DLB at Google, where it facilitates large-scale Generative AI inference for thousands of different machine learning models and millions of requests per second, we detail the design choices and practical experiences gained from the system in production. Analysis of production migrations demonstrates that DLB yields statistically significant latency reductions compared to the legacy baseline, including a 17\% decrease in median latency and a 13\% decrease at the p95 tail.

cs.DC↗
arXiv2026

Stellar Colosseum: A Many-Agent Harness for Long-Horizon Research in Mathematics and Theoretical Computer Science

Language models can produce plausible short proofs, but may still be unreliable on long-horizon research problems, where progress depends on a sequence of uncertain and interdependent decisions. We introduce Stellar Colosseum, a model-agnostic harness for allocating inference across research in mathematics and theoretical computer science. Colosseum explores alternative strategies before proof construction, uses a readiness gate to decide when a route is mature enough to decompose, represents the proof plan as interdependent section-level subproblems, and routes verifier findings back to the affected part of the argument. Across these stages, it generates candidates in parallel, attacks them with targeted falsification, and combines candidates and their critiques into a single research artifact through overlapping random-sample tree aggregation. The Colosseum workflow has been integrated into Google Antigravity's Teamwork framework as the Long Proof pattern. We demonstrate the capabilities of Colosseum through open-ended research and evaluations on theorem-proving and competitive programming benchmarks. Using Colosseum with Gemini 3.1 Pro, we obtain several new results that address open problems arising from papers published at top venues such as FOCS and JMLR. On TCS-Bench, a benchmark of research-level theorem-proving tasks drawn from papers published at FOCS, STOC, and SODA, Colosseum achieves 71.0% accuracy using Gemini 3.1 Pro and Gemini 3.7 Flash. In a separate Codeforces evaluation using Gemini 3.1 Pro, the proof-oriented pipeline with execution feedback solves 218 of 222 problems.

cs.AI↗
arXiv2026

Causal inference with bipartite designs: A generalized propensity score approach

Bipartite experiments, in which one set of units receives a treatment while outcomes are measured on another set, have recently garnered attention for their ability to capture interference across two distinct populations, such as buyers and sellers in online marketplaces. However, analyzing these experiments can be challenging, given that exposure is neither purely binary nor independent across units. In this paper, we propose a unified framework for causal inference in bipartite designs that leverages generalized propensity scores (GPS) to estimate exposure-response functions. Under standard unconfoundedness assumptions, we show that our GPS-based estimators are unbiased and derive theoretical bounds on their variance. We further introduce practical modeling and weighting strategies, along with double deconfounding methods, that integrate the GPS into both the outcome model and the assignment mechanism. Through extensive simulations, we demonstrate that these approaches achieve substantial bias reduction compared to naive methods. We also illustrate their effectiveness in a real-world application using an Amazon Pet Supplies dataset, where controlling for network structure proves critical to drawing valid causal conclusions. Our results underscore the importance of bipartite designs in contexts with significant interference and highlight how GPS-based methods can bolster the reliability of causal effect estimates in such settings.

stat.ME↗
arXiv2026

TCS-BENCH: Benchmarking State-of-the-Art Generative AI Theoretical Computer Science Research Ability

We introduce TCS-Bench, a benchmark for evaluating Large Language Models (LLMs) on research-level Theoretical Computer Science (TCS) proof generation. TCS-Bench consists of theorem-proving tasks from papers published at top theoretical computer science venues (STOC, FOCS, and SODA). Each task provides the necessary context to derive a self-contained proof for a target result. We evaluate state-of-the-art models on this benchmark. We verify the correctness of generated proofs via a verification agent, and further benchmark the verifier against human-expert proof judgements on a set of target statements and generated proofs pairs. Our reference verifier achieves over 90% accuracy on the expert labeled set.

cs.CL↗
arXiv2026

Nearly Optimal Strong Coresets for $\ell_p$ Subspace Approximation

We study strong coresets for $\ell_p$ subspace approximation. Given a matrix $A\in\mathbb{R}^{n\times d}$, the goal is to sample and rescale a small number of its rows to obtain $SA$ such that $\left\|SA(I-P_F)\right\|_{p,2}^p=(1\pm\varepsilon)\left\|A(I-P_F)\right\|_{p,2}^p$ simultaneously for every subspace $F\subseteq\mathbb{R}^d$ of dimension at most $k$, where $P_F$ is the orthogonal projector onto $F$. Woodruff and Yasuda [WY25] (FOCS 2025) obtained coreset sizes $\widetilde{O}_p(k\varepsilon^{-4/p})$ for $1\leq p<2$ and $\widetilde{O}_p(k^{p/2}\varepsilon^{-p})$ for $p>2$. We improve these bounds to $\widetilde{O}_p(k\varepsilon^{-2})$ and $\widetilde{O}_p(k^{p/2}\varepsilon^{-2})$, respectively. For $1\leq p<2$, our algorithm runs in $\widetilde{O}_p(\mathrm{nnz}(A)+d^ω+k\varepsilon^{-2})$ time. The resulting coreset size matches the sampling lower bound [LWW21] up to logarithmic factors when $k+1\geq C\log(1/\varepsilon)$ for an absolute constant $C$. For $p>2$, our algorithm runs in $\widetilde{O}_p(\mathrm{nnz}(A)+d^ω)$ time, matching the running time of the framework of Woodruff and Yasuda. We use different techniques in the two regimes. For $1\leq p<2$, we combine a bicriteria low-rank split with Lewis-weight sampling and empirical-process bounds independent of the output dimension. For $p>2$, we give a sharper analysis of the Woodruff-Yasuda construction. By retaining the truncation in its sampling probabilities throughout the row-count recurrence, we show that it achieves the improved $\varepsilon^{-2}$ dependence.

cs.DS↗
arXiv2026

Proteus: Incremental Memory Activation for Long-Context Sequence Modeling

The quadratic cost of attention-based sequence models for long contexts has motivated a growing line of research on memory-based models that can compress context into a compact state. However, most existing memory models expose a static memory throughout the entire sequence. Because early tokens face no compression pressure, they occupy too many degrees of freedom and "pollute" the memory state, leaving little capacity for later context and increasing interference between what is stored and what arrives next. We study a new paradigm of incremental memory activation, where the effective capacity of memory is progressively expanded as the context grows. Imposing an early bottleneck forces the model to compress history more effectively, while unlocking fresh capacity over time reduces interference and improves retention of later context. We instantiate this paradigm in Proteus, a straightforward mechanism that can be incorporated into a broad class of neural memory architectures at no additional cost. We apply Proteus to state-of-the-art models, including SWLA, Comba, Titans, and Hope-Attention, and observe consistent improvements on standard language modeling and reasoning, as well as on long-context retrieval and understanding, with gains that grow at longer context lengths. Overall, our results show that static memory is suboptimal and that scheduling effective capacity is a simple and broadly applicable tool for sequence modeling.

cs.LG↗
arXiv2026

SSTQ:Privacy-Preserving Vector Quantization via Subsampled Stochastic TurboQuant

Achieving local differential privacy in distributed optimization while maintaining low communication cost remains challenging. Existing vector quantization methods, such as vqSGD, use high-dimensional geometric constructions but incur unfavorable dimension-dependent variance. In this work, we propose Subsampled Stochastic TurboQuant (SSTQ), a framework that combines overcomplete equal-norm tight frames, coordinate subsampling, and privacy-aware one-dimensional quantization. SSTQ includes two variants: a Flat Randomized Response version and a Metric-Aware Laplace version, the latter being better suited to higher codebook bit-width regimes. We show that SSTQ achieves optimal mean squared error scaling while using only $\lceil \log_2 N \rceil + b$ bits per client, where $N = Θ(d)$ is the frame size. We also derive a surrogate privacy-aware codebook objective that reduces the codebook-dependent MSE scaling from $O(4^b)$ to $O(2^b)$. Finally, we empirically evaluate SSTQ against established baselines on federated learning tasks using CIFAR-10 and Fashion-MNIST, demonstrating favorable utility and communication efficiency. Some of the analytical derivations were first obtained using a fully automated Gemini-based agentic system developed internally at Google. The authors have verified those derivations and edited them for clarity of presentation.

cs.LG↗
arXiv2026

A Near-Optimal Lower Bound for Prefix-Matrix Factorizations

For the $n\times n$ lower-triangular all-ones matrix $Q$, we prove a near-optimal lower bound \[ γ_{2,1}(Q) := \inf_{Q=AB} \|A\|_{2\to\infty}\|B\|_{1\to1} = Ω\!\left( \frac{\log^{3/2}n}{(\log\log n)^{3/2}} \right), \] where the infimum ranges over real factorizations of arbitrary finite inner dimension. This cost is a central parameter in space bounds for factorization-based rank and quantile estimation in turnstile streams and in error bounds for matrix mechanisms for continual counting under pure differential privacy. The proof combines right-sided Haar projections with a scale-dependent numerical-sparsity decomposition of the rows of $B$. At each scale, a rank--Frobenius argument shows that the numerically sparse rows cannot account for all of the required Schatten $2/3$ mass, while a Haar projection estimate bounds the contribution of the remaining rows. Summing these bounds over the dyadic scales yields the result. The proof was obtained using a fully automated Gemini-based agentic system developed internally at Google. The authors verified the proof and made minor revisions.

cs.DS↗
arXiv2026

Pairwise-Independent Dithering for Single-Stage Hadamard Quantization

Quantizing high-dimensional vectors is fundamental to similarity search, distributed learning, and model compression. Feng, Indyk, Kapralov, Krachun, and Prokhorov established sharp guarantees for an unbiased dithered quantizer based on a randomized Hadamard transform [FIK+26]. Their $1/d$-scale inner-product estimator, however, uses a second randomized transform and residual quantization, increasing both communication and the leading constant in the proved bound. We show that this extra stage is unnecessary: pairwise-independent dithers across Hadamard coordinates suffice. The resulting unbiased single-stage estimator uses $b$ bits per coordinate and achieves \[ \mathbb{E}\!\left[ \left|\left\langle y,\widehat{x}-x\right\rangle\right|^2 \right] \leq \left(\frac{3π\sqrt{3}}{2}+o(1)\right) \frac{\lVert y\rVert_2^2}{d\,4^b}, \] as $b\to\infty$, with a dimension-free $o(1)$ term uniform over unit inputs and fixed queries. Compared with the two-stage construction of Feng et al., it eliminates the residual-stage $O(d)$-bit payload and reduces the leading upper-bound constant by a factor of approximately $5.93$. The proof was first obtained using a fully automated Gemini-based agentic system developed internally at Google. The authors have verified the proof and edited it for clarity of presentation.

cs.DS↗
arXiv2026

The Condition-Number Barrier in Sparse Least Squares

In [AS21], Axiotis and Sviridenko conjectured that the linear dependence on the restricted condition number in sparse convex optimization cannot be improved by a polynomial-time algorithm. We establish their conjectured lower bound for least-squares objectives, conditional on the randomized exact-volume Small-Set Expansion Hypothesis in the weighted regular-graph formulation of Raghavendra, Steurer, and Tulsiani [RST12]. Concretely, for every fixed $γ\in(0,1]$, there is no randomized polynomial-time algorithm that, with probability at least $2/3$, returns a vector $x$ such that, writing $s=\lVert x\rVert_0$, \[ \lVert Ax-b\rVert_2^2 \leq \min_{\lVert z\rVert_0\leq k}\lVert Az-b\rVert_2^2+\varepsilon \quad\text{and}\quad s=O\!\left(k\,κ_{s+k}^{\,1-γ}\right), \] where $κ_r$ is the restricted condition number at sparsity level $r$. The result holds even on rational instances with $A$ of full column rank. The proof was first obtained using a fully automated Gemini-based agentic system developed internally at Google. The authors have verified the proof and edited it for clarity of presentation.

cs.DS↗
arXiv2026

Near-Optimal Dimension Lower Bounds for Single-Vector Embeddings of Maximum Inner Product Similarity

Multi-vector embeddings represent items by point clouds and compare query and document point clouds using Chamfer similarity, whereas single-vector embeddings use ordinary inner products. For singleton queries, Chamfer becomes maximum inner product similarity (MAX-IP). In our setting, MUVERA gives dimension $m^{O(1/ε^2)}$ [DHJ+24], whereas the previous lower bound $(ε^2m)^{Ω(1/ε)}$ [Jay26] left a gap between $1/ε$ and $1/ε^2$ in the exponent of $m$. We nearly close this gap. For every fixed $δ\in(0,1)$, there are constants $A_δ,c_δ>0$ such that, for all sufficiently small $ε>0$ and every $m\ge(1/ε)^{A_δ}$, there exist unit query vectors and document point clouds of at most $m$ unit vectors for which every single-vector approximation of all pairwise MAX-IP values to additive error $ε$ has dimension \[ D \ge m^{c_δ/ε^{2-2δ}}. \] This holds even for fully data-dependent representations chosen after seeing the dataset. It also applies to Chamfer because all queries are singletons. Since $δ$ can be arbitrarily small, the exponent approaches the $O(1/ε^2)$ dependence of the upper bound. The proof combines Sherstov's pattern matrix method with polynomial-size, constant-width DNF formulas computing functions of approximate degree $Ω(k^{1-δ})$. Uniform-width padding and a block encoding create an $Ω(ε)$ gap. A dummy coordinate then equalizes all false inputs, yielding a unit-sphere MAX-IP matrix that is an exact two-valued affine image of the DNF pattern matrix with gap at least $8ε$. This allows the approximate-rank bound to apply. The proof was first obtained using a fully automated Gemini-based agentic system developed internally at Google. The authors have verified the proof and edited it for clarity of presentation.

cs.DS↗
arXiv2026

Language Models Need Sleep: Learning to Self-Modify and Consolidate Memories

The past few decades have witnessed significant advances in the design of machine learning algorithms, from early studies on task-specific shallow models to more general deep Large Language Models (LLMs). Despite showing promising results in tasks that require instant prediction or in-context learning, existing models lack the ability to continually learn and effectively transfer their temporal in-context knowledge to their long-term parameters. Inspired by human learning process, we introduce a ''Sleep'' paradigm that allows the models to continually learn, distill their short-term fragile memories into stable long-term knowledge with replay, and recursively improve themselves with ''Dreaming'' process. In more detail, sleep consists of two stages: (1) Memory Consolidation: an upward distillation process, called Knowledge Seeding, where the memories of a smaller-self are distilled into a larger network to provide more capacity while preserving the knowledge. As a proof of concept, we present a new Generalized Distillation process for {Knowledge Seeding} (i.e., the combination of on-policy distillation with Reinforcement Learning (RL)-based imitation learning); (2) Dreaming: a self-improvement phase, where the model uses RL to generate a curriculum of synthetic data to rehearse new knowledge and refine existing capabilities without human supervision. Our experiments on long-horizon, continual learning, knowledge incorporation, and few-shot generalization tasks support the importance of the sleep stage.

cs.LG↗
arXiv2026

Quota Marketplace: Dynamic Pricing for Efficient Allocation of ML Training Resources

The escalating demand for Machine Learning (ML) training resources in recent years has resulted in a substantial gap between the high demand and the available supply. Efficient allocation of these scarce and expensive resources is crucial for organizations to maximize their return on investment. Existing resource allocation mechanisms, like Karma [OSDI'23], are designed to guarantee Pareto efficiency and max-min fairness in settings with dynamic (time-varying) user demands, but fail to preserve these key properties in the presence of demands with heterogeneous values. Given the ubiquity and inevitability of heterogeneity in organizational values of different workloads, effective resource allocation policies must accommodate these variations. In this paper, we describe the design, implementation, deployment, and theoretical analysis of Quota Marketplace, a market-based mechanism to efficiently allocate ML training chips (like GPUs), explicitly addressing scenarios with demands of heterogeneous value. We detail the implementation of this mechanism within Google and present metrics that demonstrate its impact. We also discuss many business-critical requirements that the Quota Marketplace handles quite effectively, and document the gains and opportunities it has unlocked. We establish theoretically how this market-based approach achieves the essential properties of Pareto efficiency and max-min fairness by allowing the users to express the value of their workloads and enabling dynamic resource pricing based on supply and demand fluctuations. Ultimately, the market facilitates resource allocation that aligns with organizational priorities.

cs.LG↗
arXiv2026

HyperVAttention: Efficient Sparse Attention with Spatio-Temporal Clustering for Video Diffusion

Video Diffusion Transformers (VDiTs) have demonstrated significant capabilities in high-fidelity video generation. However, their ability to produce long-duration videos is fundamentally constrained by the quadratic complexity of the self-attention mechanism. Recent clustering-based sparse attention methods improve the quality-speed trade-off by grouping semantically similar tokens, but their practical efficiency remains limited by two bottlenecks: substantial clustering overhead and low CTA utilization caused by irregular cluster-induced blocks. We propose HyperVAttention (HVA), a training-free sparse attention framework that addresses both bottlenecks jointly. To reduce clustering overhead, we introduce 3D local-window clustering, which exploits the spatio-temporal locality of video tokens to restrict centroid search to fixed local neighborhoods, and implement it with a custom Triton kernel for efficient execution. We further propose a hybrid clustering strategy that performs full clustering only at anchor steps and updates only subset tokens at intermediate steps, leveraging the temporal stability of cluster assignments across denoising steps. To improve CTA utilization, we present hardware-aware cluster merging that minimizes CTA-aligned execution cost through parallel agglomerative merging, improving block density and approximation fidelity by utilizing idle tile capacity. Together, these components reduce clustering overhead, avoid redundant updates, and better align sparse attention with the fixed tile structure of modern GPU kernels. Experiments on Text-to-Video generation show that HVA establishes a new Pareto frontier for training-free sparse attention in video diffusion, reducing end-to-end latency by up to $2.13\times$ while improving fidelity over existing training-free sparse attention baselines.

cs.CV↗
arXiv2026

Geometric Signatures of Reasoning: A Spectral Perspective on Task Hardness

Chain-of-thought (CoT) reasoning enables large language models (LLMs) to solve complex problems by generating intermediate reasoning steps. While much attention has been paid to the length and content of these reasoning chains, far less is known about their internal geometry. We study the \emph{geometry} of CoT trajectories in the hidden state space of transformer models, formalizing each reasoning chain as a discrete curve in $\mathbb{R}^d$ and characterizing it through spectral, positional, and kinematic geometric functionals. We introduce the effective dimension $d_ρ$ as a measure of trajectory complexity and show theoretically that trajectories with flatter eigenvalue spectra correspond to harder tasks, as they explore more of the hidden dimensions. Lastly, we explore how kinematic features of the trajectory, mean position, positional dispersion, initial and current hidden states, mean velocity, mean speed, and speed dispersion, can be used to predict solution correctness before generation is complete, and may inform future early-stopping strategies. Experimentally, on mathematical reasoning problems from the MATH500 dataset, $d_ρ$ achieves $0.93$ AUC in distinguishing easy from hard problems, while kinematic features potentially can predict correctness from only the first $20\%$ of generated tokens. These correctness signatures transfer across questions of varying difficulty, establishing that the shape of a model's internal reasoning trajectory is a principled window into both task hardness and solution quality.

cs.LG↗
arXiv2026

Towards Automating Scientific Review with Google's Paper Assistant Tool

Artificial intelligence is driving a revolution in scientific discovery, accelerating everything from hypothesis generation to mathematical theorem proving. However, this rapid acceleration is creating a systemic challenge: traditional human peer review cannot scale to match the influx of AI-assisted science. Ultimately, to resolve this tension, we must also deploy AI to accelerate the verification and review process itself. To frame the discussion around this transition, we propose a taxonomy consisting of four progressive levels of AI-human collaboration in scientific evaluation, and discuss various trade-offs involved with each. As a step toward this future, we introduce the Paper Assistant Tool (PAT), an agentic AI framework built for deep scientific review and verification. PAT ingests full scientific manuscripts and produces a comprehensive evaluation, checking theoretical results, validating experiments, suggesting improvements, and identifying potential flaws. By utilizing inference scaling techniques, PAT is able to identify deeper issues than a single model call alone, achieving a 34% improvement over zero-shot recall on mathematical errors in the SPOT benchmark. Pilot deployments of PAT as a pre-submission tool for authors at two major Computer Science conferences -- STOC and ICML -- demonstrate its ability to identify critical errors and suggest substantive improvements to research papers. By catching errors early, PAT eases the cognitive burden placed on referees, while preserving their control over the outcomes of the review process.

cs.LG↗
arXiv2026

Distributed Load Balancing with Workload-Dependent Service Rates

Modern service systems, including cloud platforms and large language model inference endpoints, must distribute jobs across servers whose processing speeds depend on current workloads. At scale, centralized coordination is costly, while naive distributed policies can perform arbitrarily poorly. We study how to design a simple distributed load balancing policy that achieves globally optimal latency performance in such settings. We model the system as a bipartite queueing network with an arbitrary compatibility graph and servers with concave, workload-dependent service rates. We propose the Greatest Marginal Service Rate (GMSR) policy, which routes jobs to a connected server where it has the largest marginal impact on service rate. In a discrete-time stochastic model, we show that as time discretization is refined (shrinking time step and job size proportionally), the scaled workload process converges almost surely to a fluid limit governed by a differential inclusion. In the fluid regime, GMSR reaches an $ε$-suboptimal solution in $\mathcal{O}(δ+ \log(1/ε))$ time from any $δ$-suboptimal initial state, implying global convergence to the centrally optimal routing. When the system is overloaded, GMSR maximizes throughput, maximizes the number of stabilized backends among throughput-optimal policies, and minimizes total workload over those stabilized backends. GMSR yields a practical routing rule that requires neither demand-rate knowledge nor centralized coordination. By relying only on local information, service providers can achieve near-optimal latency performance through decentralized decisions, making the policy well suited to large-scale cloud computing, LLM serving, and other distributed service environments where centralized control is costly or infeasible.

cs.DC↗