The $\mathcal{O}(N^2)$ complexity of attention over $N$ tokens remains a computational bottleneck in transformer models. Vector-Quantized (VQ) attention reduces this to $\mathcal{O}(MN)$ by representing keys with $M$ codewords, but applies uniform codebook capacity regardless of where attention mass concentrates: high-attention regions of key space may be coarsely approximated while low-attention regions waste representational capacity. We propose Adaptive Vector-Quantized (AVQ) Attention, which adaptively allocates codebook capacity based on attention importance. Starting from a small set of codewords, our method identifies the most important codes during the forward pass and refines them with pre-learned child codewords, achieving fine-grained quantization where it matters most while maintaining coarse quantization elsewhere. We develop an implementation using custom Triton kernels that enables the full adaptive refinement process, including importance scoring, child codeword insertion, and parent contribution replacement, to be carried out within the tiled computation paradigm of Flash Attention with minimal overhead. Our approach maintains $\mathcal{O}(MN)$ complexity while achieving improved accuracy-efficiency trade-offs compared to fixed-codebook VQ-attention.
W.V.S.O. van den Dool, Patrick Forré, A. Habibian et al.· 0 citations
For a Markov kernel $T$ with an invariant probability measure $\pi$, we give a self-contained proof of the Markov chain convergence theorem via a criterion called asymptotic equivalence with the target. It assumes two parts about the Lebesgue decompositions of $T^n_x$ and $\pi$ for every starting point $x$: 1.) asymptotic absolute continuity: the singular mass $\mathrm{sing}(T^n_x \mid \pi)$ tends to $0$; and, 2.) asymptotic domination of the target: the singular mass $\mathrm{sing}(\pi \mid T^n_x)$ tends to $0$, as $n \to \infty$. Assuming a jointly measurable density for the absolutely continuous part of each iterate $T^n$ w.r.t. $\pi$, this criterion is sufficient and necessary for convergence. A positive minorant density version of it is verified in three cases: i.) $T$ has a positive transition density w.r.t. $\pi$; ii.) $T$ consists of an absolutely continuous part with positive transition density together with an atom at the starting point, which covers the Metropolis-Hastings algorithm; iii.) the transition density is positive only after a finite number of steps that may depend on the starting point $x$. To demonstrate our general criterion, we investigate the Gibbs sampler with random scan and the parallel tempering algorithm. Furthermore, we show that in all mentioned settings Birkhoff's ergodic theorem applies, so as to obtain the strong law of large numbers. Throughout this paper, neither irreducibility, nor aperiodicity, nor recurrence, nor couplings, nor splitting constructions, nor small sets are used. In all results, the state space is a general measurable space with no structure beyond a $\sigma$-algebra. That joint measurability is assumed of the Markov kernel, not of the space; countable generation supplies it. None of the theorems proved here is new; what is offered is a short route to a single, widely applicable Markov chain convergence criterion.