RetoMaton: Neural Retrieval Automata for Language Models

• RetoMaton is a retrieval automaton that restructures a language model's datastore into a weighted finite automaton through pointers and clustering.
• It mitigates the high cost of per-token kNN search by traversing clustered states representing sequential context and corpus continuations.
• Empirical results show it can save up to 83% of kNN searches while reducing perplexity and improving domain adaptation across benchmarks.

RetoMaton is a retrieval automaton for retrieval-augmented language modeling that approximates repeated datastore search by organizing a standard retrieval datastore as a weighted finite automaton (WFA) rather than a flat list. In its original formulation, it is an unsupervised structure built on top of the datastore used by models such as kNN-LM, with the explicit aim of avoiding an expensive nearest-neighbor search at every decoding step while preserving the standard interpolation of a pretrained LLM with a nonparametric retrieval distribution (Alon et al., 2022). Subsequent work generalizes the same automaton-centered idea into a broader neuro-symbolic retrieval framework for frozen LLMs, including a local, task-adaptive WFA intended to improve traceability, modularity, and domain alignment without parameter updates (Mamidala et al., 22 Aug 2025).

1. Origin, problem setting, and core intuition

RetoMaton was introduced in the context of retrieval-based LLMs, where a base autoregressive LLM is augmented at test time with examples retrieved from an external datastore. The bottleneck it addresses is the inference cost of per-token nearest-neighbor search: in models like kNN-LM, datastore lookup is typically performed for every next-token prediction, and that search is much slower than the base LM forward pass (Alon et al., 2022).

The central intuition is sequential. If, at time $t$, the model retrieves datastore entries whose contexts are close to the current test context, then the successors of those entries in the corpus are good candidates for what will be relevant at time $t+1$. RetoMaton captures that observation by storing pointers between consecutive datastore entries and then clustering similar entries into states so that related contexts can share outgoing continuations. This turns repeated flat retrieval into traversal of a symbolic transition system whose states are clusters of datastore entries and whose edges reflect observed corpus succession patterns (Alon et al., 2022).

This design preserves the overall retrieval-based modeling recipe of kNN-LM—pretrained LM, hidden-state keys, next-token values, distance weighting, and interpolation through $\lambda$—while changing the datastore representation and inference procedure. A plausible implication is that the method is best understood not as a replacement for retrieval, but as a structured reparameterization of retrieval over sequential corpus memory.

2. Formal construction of the retrieval automaton

The original formulation starts from the standard retrieval-based LM setup. Given a context sequence $c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$, the base LM defines a next-token distribution $p_{LM}(w \mid c^{(t)})$. If $f$ maps a context into a fixed-dimensional hidden representation, the datastore over a text collection $\mathcal{D}$ is

$$\left(\mathcal{K},\mathcal{V}\right) = \left\{ \left(f\left(c_i\right),w_i\right) \mid \left(c_i,w_i\right)\in \mathcal{D} \right\}.$$

Standard kNN-LM retrieves the $k$-nearest neighbors $\mathcal{N}$ of $t+1$0 under a distance $t+1$1 and defines

$t+1$2

with interpolation

$t+1$3

RetoMaton enriches each datastore entry by adding a pointer:

$t+1$4

where $t+1$5 points to the next datastore entry in the text collection. Its modeling assumption is

$t+1$6

By itself, this mechanism only replays exact corpus continuations, so the second ingredient is clustering entries into automaton states. If $t+1$7 clusters are used, the automaton has $t+1$8 non-initial states. The automaton is

$t+1$9

where $\lambda$0 is the state set, $\lambda$1 is the vocabulary, $\lambda$2 is an initial state, $\lambda$3 is the transition function, and $\lambda$4 is a transition-weight function. Two properties are essential: transitions may go to a set of states, and transition weights are dynamic, because $\lambda$5 depends on the current context representation (Alon et al., 2022).

Let $\lambda$6 map each datastore entry to its cluster, $\lambda$7 denote the entries in state $\lambda$8, and $\lambda$9 dereference pointers. Then the transition set is

$c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$0

In effect, the datastore becomes a symbolic graph of clustered contexts and successor relations rather than a collection of independent entries.

3. Inference, traversal, and retrieval distribution

At inference time, RetoMaton runs in parallel with the base LM. It does not eliminate exact retrieval altogether; instead, it alternates between occasional full kNN search and cheaper automaton traversal. A traversal starts with a full search producing nearest neighbors $c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$1 for $c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$2, and the active state set is initialized as the union of the states containing those neighbors:

$c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$3

After the model generates token $c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$4, the next active-state candidate set is

$c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$5

The decision to continue traversal or restart with a fresh search is governed by a threshold $c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$6. The update rule is

$c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$7

The appendix notes the extremes: $c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$8 uses traversal whenever at least one next state exists, while $c^{(t)} = (w^{(1)}, \ldots, w^{(t-1)})$9 forces kNN search at every step. This makes $p_{LM}(w \mid c^{(t)})$0 the principal quality–efficiency control.

RetoMaton also defines a retrieval-style token distribution over active states. The state weight is

$p_{LM}(w \mid c^{(t)})$1

and the automaton distribution is

$p_{LM}(w \mid c^{(t)})$2

Interpolation with the base LM is then

$p_{LM}(w \mid c^{(t)})$3

A common misunderstanding is to treat RetoMaton as a search-skipping heuristic that simply backs off to the base LM. The formalism shows otherwise: when it skips an exact kNN lookup, it still computes a retrieval-augmented distribution $p_{LM}(w \mid c^{(t)})$4 rather than reverting to $p_{LM}(w \mid c^{(t)})$5 alone (Alon et al., 2022).

4. Empirical behavior in language modeling and domain adaptation

The original evaluation covers in-domain language modeling and domain adaptation. On WikiText-103, using a 103M-token training set and a 247M-parameter Transformer LM from the original kNN-LM work, the datastore contains 103M entries and is clustered with $p_{LM}(w \mid c^{(t)})$6M. In this setting, RetoMaton can save 81% of kNN searches while matching the perplexity of standard kNN-LM; the abstract states the broader result as saving up to 83% of nearest-neighbor searches over kNN-LM without hurting perplexity. At FoSS $p_{LM}(w \mid c^{(t)})$7, perplexity drops from 16.65 for kNN-LM and 16.35 for the Adaptive Retrieval baseline to 16.08 for RetoMaton, a reduction of 1.85 perplexity relative to standard kNN-LM (Alon et al., 2022).

For Law-MT domain adaptation, the datastore and automaton are built from 19M law-domain tokens while the base LM is a large Transformer pretrained on WMT News Crawl. Here $p_{LM}(w \mid c^{(t)})$8K, giving average cluster size around 100. At FoSS $p_{LM}(w \mid c^{(t)})$9, perplexity falls from 12.34 for kNN-LM and 12.01 for Adaptive Retrieval to 10.49 for RetoMaton. The paper reports that the gains are stronger than in WikiText-103 and attributes this partly to higher n-gram repetitiveness in Law-MT: 62% of 5-grams in Law-MT validation appeared in training, versus 21% in WikiText-103. This suggests that pointer-following and clustered continuation sharing are particularly effective when local continuation structure is repetitive (Alon et al., 2022).

The paper also tests a fine-tuned target-domain LM. Starting from a Law-MT fine-tuned LM with perplexity 8.61, standard kNN-LM gets 7.93, Adaptive Retrieval gets 7.81, and RetoMaton gets 7.10 at FoSS $f$0, a 17.5% relative reduction over the fine-tuned LM. At FoSS $f$1, RetoMaton still gets 7.15, nearly identical to its FoSS $f$2 result.

The qualitative analysis is consistent with the mechanism. RetoMaton often predicts multi-token sequences without any fresh search, and the paper reports that 98% of validation tokens in WikiText-103 belong to n-grams of length $f$3 that either started or continued an automaton traversal. A plausible implication is that the automaton is exploiting both exact repeated fragments and shared continuation structure rather than functioning purely as a memorizer.

5. Ablations, implementation details, and efficiency trade-offs

The ablation study isolates the two main ingredients: pointers and clustering. A w/o clustering variant, which uses only pointers, already improves substantially. On WikiText-103 it reaches 16.12 perplexity at FoSS $f$4, better than the baselines, and can save more than 60% of searches while matching kNN-LM perplexity. The authors conclude that pointers provide most of the gains at lower FoSS, while clustering becomes especially important at high FoSS because it enables longer stretches of search-free traversal through shared continuations and compositional generalization (Alon et al., 2022).

Clustering granularity materially affects behavior. On WikiText-103, 500K and 1M $f$5-means clusters perform similarly and clearly better than 100K, which is too coarse. On Law-MT, 400K means performs best at FoSS $f$6, but becomes worse than coarser alternatives at very high FoSS. The paper also evaluates a cheaper greedy clustering algorithm from He et al. (2021): it can perform well at FoSS $f$7 but degrades more as FoSS rises, likely because clusters are smaller and support less durable traversal. This makes state formation an explicit quality–efficiency trade-off.

Implementation-wise, the method is built on the original kNN-LM code and uses FAISS both for nearest-neighbor search and for one-time clustering. It uses the same baseline hyperparameters, including $f$8 nearest neighbors for full search, and stores keys in fp16. During traversal, the number of entries used in computing $f$9 is capped at $\mathcal{D}$0 total entries across active states, preferring directly pointed-to entries and otherwise sampling cluster members randomly. This keeps scoring manageable and aligned in scale with the full-search baseline (Alon et al., 2022).

Efficiency is reported primarily through FoSS—fraction of saved searches—rather than wall-clock time, because FoSS is hardware-independent and reproducible. Supplementary analysis shows wall-clock savings track FoSS closely up to an additive constant dependent on hardware and implementation. With a CPU index, saved wall-clock time is nearly on the ideal $\mathcal{D}$1 line; with GPU index and clustering there is a fixed overhead, but significant benefits begin around FoSS $\mathcal{D}$2.

6. Later reinterpretation as local neuro-symbolic memory

A later paper, "Rethinking Reasoning in LLMs: Neuro-Symbolic Local RetoMaton Beyond ICL and CoT," repositions RetoMaton as a more general neuro-symbolic retrieval framework for LLMs in which a frozen LLM is augmented with symbolic memory organized as a WFA (Mamidala et al., 22 Aug 2025). In that treatment, RetoMaton is contrasted with prompt-based reasoning methods such as In-Context Learning (ICL) and Chain-of-Thought (CoT), which are described as sensitive to prompt structure, formatting, and example ordering, or as yielding fragile intermediate steps.

The main extension is local RetoMaton, which replaces a global datastore with a local, task-adaptive WFA built directly from external domain/task-relevant corpora. The paper distinguishes three memory-construction regimes: a global datastore built from a broad corpus such as WikiText-103, a domain-aligned datastore built from task-relevant corpora, and a local datastore built from the most task-specific material—for example, the MMLU training split, the GSM8K training split, or query-specific evidence documents for TriviaQA. The claim is that locality reduces retrieval noise, improves alignment with the model’s latent predictive structure, lowers inference overhead by using much smaller datastores, and preserves symbolic traceability (Mamidala et al., 22 Aug 2025).

The formalism again uses hidden-state clustering and transition-constrained retrieval. In the appendix, the fallback hierarchy is made explicit as automaton-constrained $\mathcal{D}$3 cluster-based $\mathcal{D}$4 global $\mathcal{D}$5 kNN. The paper also notes a terminological subtlety: although it speaks of “deterministic” or “structured” transitions, its formal transition function is explicitly non-deterministic,

$\mathcal{D}$6

so “deterministic” is used in the sense of traceable and rule-governed retrieval behavior, not strict DFA determinism. This addresses a possible misconception about the symbolic component (Mamidala et al., 22 Aug 2025).

Empirically, the paper reports average gains of 4.48\% with LLaMA and 2.78\% with Gemma across TriviaQA, GSM8K, and MMLU. For LLaMA, it reports 41.96 EM and 48.23 F1 on TriviaQA, 30.54 on MMLU, and 50.49 on GSM8K for Local RetoMaton, compared with 39.62/45.51, 23.00, and 48.75 for the baseline model. On GSM8K, it also reports improved calibration-style metrics for Local RetoMaton: PPL = 2.7787, KLD = 0.0359, and NLL = 1.0193 (Mamidala et al., 22 Aug 2025).

7. Limitations, scope, and relation to neighboring methods

The original RetoMaton paper presents the method as an approximation to repeated kNN search, and several limitations follow from that fact. If traversal drifts away from the truly nearest neighbors, quality may drop; $\mathcal{D}$7 mediates that trade-off. Clustering quality matters substantially: clusters that are too coarse are noisy, while clusters that are too fine reduce sharing and shorten search-free runs. Traversal is not free, because scoring active states and retrieving their member entries adds overhead, especially with clustering on GPU. The method also works best when future continuations are predictable from retrieved neighbors, which the paper states is more true in repetitive domains like law than in noisier open-domain text. Finally, although no LM retraining is needed, datastore construction and clustering for a very large corpus remain a significant preprocessing step (Alon et al., 2022).

Relative to neighboring methods, what remains the same as in kNN-LM is the retrieval-and-interpolation philosophy. What changes is the representation of retrieval memory as a pointer-and-cluster-based automaton. Relative to Adaptive Retrieval, the decisive distinction is that RetoMaton can skip exact search while still producing a retrieval-augmented distribution rather than backing off entirely to the base LM. Relative to broader systems such as RAG, the later local-RetoMaton work emphasizes that retrieval occurs over hidden representations organized by automaton transitions rather than over raw retrieved text inserted into the prompt (Mamidala et al., 22 Aug 2025).

In concise terms, RetoMaton is a neuro-symbolic augmentation of retrieval-based language modeling that converts a flat hidden-state datastore into a WFA with successor pointers and clustered states. Its practical appeal lies in the combination of three properties stated across the literature: no extra supervision, no retraining of the LLM, and the possibility of obtaining better perplexity and many fewer searches through structured traversal of external memory (Alon et al., 2022).