Cognithesis — an instrument for curiosity

cognithesis.← Return to the engine
← Origin

ORIGIN / 04

HIDRA: Attention Is a Search Problem

An old theory paper, dusted off: replace quadratic attention with exact k-NN retrieval, and let a router decide where to look. Untested, unrepented.

In this entry

The following is a theory I wrote a bit over two years ago and never published. It is untested. It may be wrong. I still believe the core move is sound, and I’d rather have it out in the world being argued with than sitting in a folder being ignored. I have kept the honest parts honest: there are no results here, only a design and a list of experiments that would kill it.

The idea in one sentence

Attention doesn’t have to compute the relationship between a query and every key in the context — it can look it up, with exact nearest-neighbor search over an externalized key–value store, and weight the winners with a Gaussian kernel.

Why this isn’t obviously wrong

Dense self-attention is quadratic in sequence length, which is the fact that launched a thousand architectures. The standard fixes fall into families:

  • Structured sparsity — decide in advance which positions matter (Longformer, BigBird).
  • Approximation — kernelize the softmax until the quadratic term goes away (Performer and friends).
  • Better engineering — FlashAttention’s IO-aware tiling, which keeps attention dense but makes it hurt less.

All three accept a premise: the full query–key relationship matrix is conceptually there, and we trim or approximate it for speed.

HIDRA drops the premise. The claim: for most queries, the keys that matter are the exact top-K nearest neighbors — and finding exact top-K in high dimensions is a solved systems problem, not a modeling problem. The attention operation becomes:

  1. Select the exact top-K neighbors of the query under cosine or L2 distance.
  2. Weight them by exp(-distance²/τ) — a Gaussian/RBF kernel — and normalize.
  3. Sum the corresponding values. Done.

When K equals the sequence length, this is softmax attention (up to temperature and normalization). The question is what breaks when K is 64 and the context is a million tokens.

There’s a clean equivalence hiding in step 2. If queries and keys are L2-normalized, ‖q − k‖² = 2 − 2(q·k), and the Gaussian weights collapse to a top-K softmax over dot products. So HIDRA is not some exotic replacement for attention — it’s the same operation, restricted to exact nearest neighbors instead of everything.

The “exact” part is the heresy. The efficient-attention literature leans on approximate methods — hashing, random features, quantization — and eats the recall loss. I wanted the opposite: exact selection over the stored keys, with all the approximation pushed into where you don’t look, not what you get wrong when you do.

The part that made me write it down

Exact top-K over a million keys is only cheap if the search kernel prunes aggressively, and the pruning trick is where this paper earned its keep in my head.

Store keys column-major — dimension j contiguous across all positions — and accumulate each candidate’s squared distance as a partial sum over dimensions. Now you can early-abandon: once a candidate’s partial distance exceeds the current K-th-best full distance, you stop reading its remaining dimensions. The candidate is provably out of the top-K; you’ve spent a few cache-line loads to know that instead of d.

Two multipliers stack on top:

  • Dimension ordering. If you permute dimensions so the discriminative ones come first, partial sums separate near-keys from far-keys faster, and abandon rates climb. This ordering can be learned — it’s a permutation, it’s differentiable-ish through the usual tricks, and it’s free at inference.
  • Pre-routing. A coarse partition of the space lets a router name a handful of shards worth searching before any distance is computed at all. The rest of the store is skipped without touching a byte.

The per-token cost of the core kernel is then O(K·d) after routing — not O(n·d). The context lives in memory-mapped shards on the host; hot shards cache on-device. Working memory per query is independent of context length, which is the whole point: practically unbounded context is a storage decision, not a compute decision.

And the search subsystem isn’t hypothetical — it’s the shape of search engine I’ve been building separately, nebula, a blazing-fast vector search engine. HIDRA is, in a sense, the question “what if attention were a customer of a search engine?”

The router

Search needs to be told where to look, so a small network — the intention router — maps the local context to an intention vector that predicts (1) which shards to probe and (2) optionally, which adapters or experts to invoke. Standard task loss, plus auxiliary terms to penalize candidates examined and prevent the router collapsing onto one shard.

The router is also the honest name for the risk. Attention patterns that are positional rather than content-based — copying, induction heads, “attend to the previous token” — are not nearest-neighbor problems, and a content-routed search will not find them. I flagged this in the paper as a worst case, and I’ll flag it harder here: this is the known failure mode of the entire content-based-sparsity family, and HIDRA does not escape it. The fix, if there is one, is that the router can learn to route positionally — shard zero is “recent context,” and the intention vector can learn to point there for copy-like tasks. But that’s a hypothesis about training dynamics, not a guarantee.

What would kill it

The paper closes with an evaluation plan, which I’m reproducing essentially verbatim, because pre-registration is the difference between a theory and a vibe:

  • Tasks: long-document QA and summarization; needle-in-a-haystack and multi-needle retrieval; repo-scale code completion; streaming anomaly detection.
  • Baselines: FlashAttention as the dense standard; Longformer/BigBird for sparsity; a BM25/embedding retrieval pipeline as the “just use RAG” control.
  • Metrics that matter: not just task scores — the search statistics. Fraction of candidates examined, fraction of dimensions visited, p95/p99 tail latency, recall@K of the selection step. If early-abandon isn’t abandoning, the whole cost model is a fiction.
  • Ablations: fixed-K vs. mass-targeted accumulation; learned vs. variance-based vs. random dimension order; router on/off; int8 vs. fp16 keys.

None of it has been run. The status of this work is theory.

Two years later

Why publish now? Because the surrounding pieces exist now, and theories get better when they’re load-bearing. The search-engine half of the design is real and fast. The kernel side of my work keeps proving that the interesting results come from ignoring the intended interface and talking to the metal. And the honest ledger of this blog means I can say what this is without dressing it up: a two-year-old idea that is either a week of experiments away from being obvious in hindsight, or already wrong for a reason everyone else discovered first.

The way to bet, historically, is that the objection I can’t answer is the real one. I can’t answer the positional-attention one yet. But the question — is attention a search problem? — has aged well, and I’d rather find out in public.

← Back to origin