State-space models (SSMs) offer efficient alternatives to Transformers for long sequences, but their fixed-size recurrent state limits capability on algorithmic tasks, such as retrieving past context. In this work, we examine how in-context retrieval operates in Transformer- and SSM-based language models and find that both rely on a Gather-and-Aggregate (G&A) mechanism: a Gather Head extracts relevant information from context, which an Aggregate Head integrates into representation. In both architectures, G&A concentrates in a few heads, forming bottlenecks even for simple retrieval. For example, disabling a single Gather or Aggregate Head in a pruned Llama-3.1-8B impairs retrieving the correct answer letter in MMLU, reducing its accuracy from 66% to 25%. Moreover, this retrieval bottleneck can obscure knowledge demands of tasks as the pruned model succeeds on MMLU with functioning G&A heads yet fails on other knowledge benchmarks. The bottleneck similarly extends to tasks where SSMs typically underperform, like GSM8K, BBH, and dialogue. We show that SSMs’ retrieval challenges manifest in these heads, creating smoother attention patterns instead of the sharp transitions effective G&A requires. Thus, the Transformer-SSM retrieval gap exists in just a few heads, rather than the entire language model. % Result 3: Analyzing Hybrid models This suggests a unified explanation for Transformer vs. SSM performance gap while showing how to merge their strengths. We find that pretrained hybrid models, where SSMs are combined with attention layers, delegate the role of Aggregate Heads to attention. Similarly, replacing a single G&A head in a pretrained SSM with an attention variant boosts retrieval and benchmark scores.

Transformer architectures have become a dominant paradigm for domains like language modeling but suffer in many inference settings due to their quadratic-time self-attention. Recently proposed subquadratic architectures, such as Mamba, have shown promise, but have been pretrained with substantially less computational resources than the strongest Transformer models. In this work, we present a method that is able to distill a pretrained Transformer architecture into alternative architectures such as state space models (SSMs). The key idea to our approach is that we can view both Transformers and SSMs as applying different forms of mixing matrices over the token sequences. We can thus progressively distill the Transformer architecture by matching different degrees of granularity in the SSM: first matching the mixing matrices themselves, then the hidden units at each block, and finally the end-to-end predictions. Our method, called MOHAWK, is able to distill a Mamba-2 variant based on the Phi-1.5 architecture (Phi-Mamba) using only 3B tokens. Despite using less than 1% of the training data typically used to train models from scratch, Phi-Mamba boasts substantially stronger performance compared to all past open-source non-Transformer models. MOHAWK allows models like SSMs to leverage computational resources invested in training Transformer-based architectures, highlighting a new avenue for building such models.

Zero knowledge proofs are one of the most influential concepts in theoretical computer science. In the seminal definition due to Goldwasser, Micali and Rackoff dating back to the 1980s, a computationally-bounded verifier interacts with a powerful but untrusted prover, with the goal of becoming convinced that the input is in some language. In addition to the usual requirements of completeness and soundness, in a zero knowledge proof, we protect the prover's knowledge: assuming the prover is honest, anything that the verifier can deduce after interacting with the prover, it could have deduced by itself. Zero knowledge proofs have found many applications within theoretical computer science and beyond, e.g., in cryptography, client-cloud computing, blockchains and cryptocurrencies, electronic voting and auctions, and in the financial industry. We define and study the notion of distributed zero knowledge proofs, reconciling the computational notion of zero-knowledge with the communication-based paradigm of distributed graph algorithms. In our setting, a network of verifiers interacts with an untrusted prover to decide some distributed language. As is usually the case in distributed graph algorithms, we assume that the verifiers have local views of the network and each only knows its neighbors. The prover, on the other hand, is assumed to know the entire network graph, as well as any input that the verifier may possess. As in the computational centralized setting, the protocol we design should protect this knowledge. In particular, due to the dual role of the underlying graph in distributed graph algorithms, serving as both the communication topology and the input to the problem, our protocol must protect the graph itself. We construct communication-efficient distributed zero knowledge proofs for two central problems: the 3-coloring problem, one of the poster children of computational zero-knowledge, and for the spanning-tree verification problem, a fundamental building block for designing graph algorithms. We also give a general scheme for converting proof labeling-schemes to distributed zero-knowledge protocols with related parameters. Our protocols combine ideas from computational complexity, distributed computing, and cryptography.