Permanent Domain Memory and Faster KV-cache: Tauformer, the Topological Transformer.tex

\documentclass[12pt,a4paper]{article}
% ===== PACKAGE IMPORTS =====
\usepackage[utf8]{inputenc}
\usepackage[english]{babel}
\usepackage[a4paper,left=17mm,top=22mm,right=17mm,bottom=15mm]{geometry}
% Math packages
\usepackage{amsmath,amsfonts,amssymb}
\usepackage{mathtools}
\usepackage{bm}
\usepackage[T1]{fontenc}
% Graphics and floats
\usepackage{graphicx}
\usepackage{float}
% Bibliography and citations
\usepackage{natbib}
% Links and references
\usepackage[colorlinks=true, urlcolor=blue, linkcolor=blue, citecolor=red]{hyperref}
\usepackage{url}
% Code listings
\usepackage{listings}
\usepackage{xcolor}
\usepackage{graphicx}
\usepackage{booktabs}
\usepackage{array}
\bibliographystyle{ieeetr}
% Configure Rust syntax highlighting
\lstdefinelanguage{Rust}{
	keywords={fn, let, mut, if, else, for, while, loop, match, enum, struct, impl, trait, pub, use, mod, crate, self, super, const, static, unsafe, extern, true, false, Some, None, Ok, Err, Vec, Box, Option, Result},
	keywordstyle=\color{blue}\bfseries,
	ndkeywords={String, str, i32, i64, f32, f64, u8, u16, u32, u64, usize, isize, bool, char},
	ndkeywordstyle=\color{purple}\bfseries,
	sensitive=true,
	comment=[l]{//},
	commentstyle=\color{gray}\itshape,
	string=[b]",
	stringstyle=\color{red},
	basicstyle=\ttfamily\footnotesize,
	breaklines=true,
	breakatwhitespace=true,
	tabsize=2,
	showstringspaces=false,
}
\lstset{
	language=Rust,
	frame=single,
	numbers=left,
	numberstyle=\tiny\color{gray},
	xleftmargin=2em,
	framexleftmargin=1.5em,
}
% ===== LISTINGS CONFIGURATION =====
\lstset{
	basicstyle=\ttfamily\small,
	commentstyle=\color{gray},
	keywordstyle=\color{blue},
	stringstyle=\color{red},
	numberstyle=\tiny\color{gray},
	numbers=left,
	stepnumber=1,
	numbersep=5pt,
	backgroundcolor=\color{gray!10},
	frame=single,
	rulecolor=\color{black!30},
	captionpos=b,
	breaklines=true,
	breakatwhitespace=false,
	tabsize=2,
	showstringspaces=false
}

\begin{document}
	
	\title{Topological Transformer: A Redesign\protect\\for Domain-Memory\protect\\and cheaper kernel operations}
	
	\author{Lorenzo Moriondo\\
		\small Independent Researcher - tuned.org.uk\\
		\small ORCID: 0000-0002-8804-2963}
	
	\date{5 January 2026}
	
	\maketitle
	
	\begin{abstract}
		Starting from previous research on topological spaces for vector search (spectral search with \textit{taumode}), I introduce in this paper the definition, concept implementation and testing of a new class of Transformers focused on boosting the attention mechanism.
		The Topological Transformer, or \textit{Tauformer}, applies all the features from current attention-based transformers (taking as baseline \textit{nanoGPT}) and redesigns the attention mechanism at its core to: \textit{(i)} allow delivering domain-specific context at the level of the attention mechanism for more domain-relevant token generation, \textit{(ii)} improve time per token by $\sim20\%$ and provide savings in KV-cache memory size by $\sim50\%$ and \textit{(iii)} set a potential pathway to improve Transformers performance on larger context windows with longer prompt lengths and high-dimensional embeddings; this is made possible by substituting inner-product with \textit{taumode}'s synthetic index-based distance, aiming to provide relevant linear gains in training and generation compared to current GPTs. 
	\end{abstract}
	
	\section{Introduction}
	
	Concepts and tools developed in my previous publication \textit{ArrowSpace: introducing Spectral Indexing for vector search} \cite{ArrowSpace:2025} have been here reused to redesign the attention mechanism forking the \textit{nanoGPT} \cite{nanoGPT:2022} implementation for the Transformer architecture \cite{vaswani_attention:2017}; in particular \textit{taumode}, a synthetic index based on the Rayleigh quotient \cite{Rayleigh:1877} to compute the distribution of energy in the network defined by the embeddings space. Examples about how to compute the \textit{taumode} distribution for any space of embeddings is available in the
	\textit{pyarrowspace} \href{https://github.com/tuned-org-uk/pyarrowspace/tree/main/tests}{repository}.
		
	\textbf{Tauformer} makes possible downstream technical improvements in computing and memory usage for the Transformer (GPT) and they are detailed in the following sections. These improvements are a consequence of pursuing the concept of providing a memory layer to LLMs that brought forward the idea of leveraging \textbf{distilled knowledge graphs (dkb) to deliver domain information to the attention mechanism at token generation level}. For dkb is intended here for example: \textit{(i)} Text embeddings generated by feeding to an embedding model (in the tested example: TSDAE \cite{wang-etal-tsdae:2021}) a representative corpus of the domain to map, this provides a vector space with manageable number of dimensions (384 in the example) that is a prerequisite to build the context windows at the attention level; \textit{(ii)} Graph embeddings of triples generated using techniques like the SEPAL workflow demonstrated in \cite{lefebvre_sepal:2025}, again the result is a vector space with a manageable dimensionality.\\
	A coherent dkb is generated starting from the vector space defined by its embeddings (and its Graph Laplacian computed using \textit{arrowspace} \cite{ArrowSpace:2025}), this is considered from the point of view of the LLM system as a persistent memory that is founded on real-world-defined relations. Assuming a well-designed embedding pipeline, text embeddings from a curated text corpus or a curated Knowledge Base turned into graph embeddings become the numerical ground truths for the topology of the context, that is eventually delivered to the attention mechanism.
	Theoretically the Tauformer can also work, as any other transformer, on images and audio vectors or a mix of them but one of the critical points focuses on the process of producing the embeddings, how the vector space is designed and how its dimensionality fits in the attention mechanism.\\The choices for: embeddings, attention mechanism and seeding the latent space are taken as design constraints for the data pipeline from the corpus to the latent space of the model.	
	What is meant with seeding the latent space? Once the attention mechanism is working with \textit{taumode} synthetic scores the content of the dkb percolates in the context windows through the usual $Q, K, V$ iterations. Each distance measurements computed for $Q, K, V$ is considered in the optics of the corpus/dataset from which the original embeddings were generated, delivering domain-specific metadata directly in the attention mechanism and through that to the latent space.
	
	\section{Statement of Need}
	
	This research is motivated by encoding at generation time an explicit (controllable and interpretable) domain structure beyond what the model can implicitly absorb into weights during training.  Tauformer keeps the familiar $Q,K,V$ and causal softmax pipeline but replaces the dot-product kernel with a distance metrics based on domain-specific manifold: each token/head is mapped to a topological signal and built into a scalar (\textit{taumode} score) derived from the Graph Laplacian (manifold); this way attention is driven by distances in that manifold-aware scalar space rather than raw geometrical vector similarity. The hypothesis is that this process makes the model more scoped to the context in the training phase (forward) and in the generating phase (decoding), because weights and attention scores can be biased toward tokens that are similar under the learned domain manifold (i.e. a knowledge graph, citation graph, text embedding or any dataset originally passed to the embedding model). By hypothesis, this has the potential of improving contextual faithfulness relative to purely dot-product-based retrieval inside the context window. Not having permanent reference about the context is especially consequential in scientific and knowledge-intensive settings, where the desired notion of "relevance" is frequently defined less by raw generic geometrical proximity and more by the topology of domain relations and how information (in \texttt{arrowspace}, energy) propagates over them.
	
	Tauformer is also motivated by the practical difficulty of scaling traditional GPTs to very long contexts and large embeddings, in particular at inference time: full-sequence self-attention has $O(T^2)$ compute in both the dot-product logits and the value aggregation, and decode-time still grows $O(T)$ per token while KV-cache memory grows linearly with $T$.  Tauformer’s cache design reduces KV memory by storing values plus a compact key-side scalar $(k_i, \lambda_i)$ instead of full $K$ and $V$ tensors, yielding roughly a halving of KV-cache memory per layer (up to a small overhead for the scalar stream).  With a sparse Laplacian from a domain manifold, the extra cost of computing the \textit{taumode} scalars can shift from $O(D^2)$ to a sparsity-dependent cost $O(\mathrm{nnz}(L))$ (non-zero elements of the Laplacian matrix), making the incremental overhead largely independent of sequence length and therefore more attractive as context windows and embeddings dimensions grow.
	
	Tauformer addresses this by redesigning attention around a synthetic, topology-aware index computed over the Graph Laplacian of the domain's embedding space. Beside accuracy and controllability, this substitution is also motivated by systems constraints: the approach reduces dependence on storing key/value histories (full vectors) for long contexts and enables more memory-efficient inference, opening space for higher-dimensional embeddings and richer internal representations under fixed hardware budgets. The Graph Laplacian (gl) is usually a DxD matrix where D is the number of dimensions that for the test dataset are defined by current SOTA embeddings models to 384. So basically a constant 384x384 sparse matrix ($nnz\sim10^4$ for a realistic sparsity of 1-8\%) allows limiting the computing cost for the iteration of \textit{Q, K, V} in the attention mechanism. The Graph Laplacian can be seen as a distillation of the domain knowledge and its generation takes a cheap pre-training step that for the test dataset is $\sim300$ seconds on a common laptop hardware for a dataset of $3*10^5$ vectors on 384 dimensions.
	
	The code for the GPT implementation is available at \cite{tauformer:2025}.
	The code used to test and collect benchmark for the KV-cache is available at \cite{taugpt_kvcache:2025}.
	
	\section{Memory Model and Algorithm}
	
	\subsection{Generics}
	\label{sec:generics}
	
	 By leveraging the Graph Laplacian matrix (\textit{gl}) built on the vector space defined by the embeddings for a given domain, it is possible to define a synthetic score ($\lambda\tau$, lambda-tau or \texttt{taumode}) that can be used to spot approximate nearest neighbours of a given vector (\textit{x}) as: $\lambda\tau = taumode(x, gl)$. In the Tauformer architecture this distance metrics has been applied to the process of computing query, keys, value vectors (\textit{Q, K, V}) in the attention mechanism; allowing the substitution of $Q^T \dot K$ with $-|\lambda_q - \lambda_k|$. This also enable saving a relevant percentage of the memory space used by the KV-cache and opens to having larger vectors and attention heads because part of the computing used by the inner-product operation can be reused to compute on higher dimensional vectors.  
     These improvements sound already a nice step forward but, from knowledge engineering perspective, they are side-effects of the main concept brought forward by the implementation overall: \textit{to develop a delivery system to bring distilled domain knowledge from the vector space (representing the knowledge graph, graph of citations, or any relations-representative structure) directly into the attention mechanism}. This to test the hypothesis that token generation can be made more context-relevant if domain-specific metadata is provided at the attention level. If this is made possible a new class of Transformers that are more or less narrowly scoped for context-specific generation in a given domain can be developed.
	
	
	\subsection{Tau Attention}
	
	In the process of causal attention using the \textit{softmax} function, attention is driven by a transcendental normalisation over dot-product logits, where each new query must form inner-products against \textit{all cached keys} (length $T$) in head dimension $D$, i.e., the decode-time kernel cost scales as $O(B\,H\,T\,D)$ for the $QK^\top$ logits plus another $O(B\,H\,T\,D)$ to aggregate values, with an additional $O(B\,H\,T)$ for masking and the softmax itself.
    Where the letters are:
	\begin{itemize}
		\item \textbf{B}: Batch size (number of sequences processed in parallel).
		\item \textbf{H}: Number of attention heads.
		\item \textbf{T}: Sequence length / number of time steps (tokens) in the context window.
		\item \textbf{D}: Head dimension (the per-head vector size, typically $D = C/H$ where $C$ is the model embedding width).
	\end{itemize}
	
	 In practice this also forces the KV-cache to store both $K$ and $V$ tensors, i.e., $2\cdot B\,H_{kv}\,T\,D$ floats per layer, because the dot products cannot be reconstructed without the full key vectors.  As context length grows, this "match against all previous keys and renormalise" pattern dominates both compute and memory, since every step repeats the same $D$-dimensional comparisons against an ever-growing set of cached vectors. Obviously this can be become a problem for high $D$.
	
	Tauformer’s \textit{taumode} mechanism replaces the dot-product kernel with a scalar spectral signature per head vector: for each query/key $x\in\mathbb{R}^D$, it computes a bounded Rayleigh quotient energy $x^\top L x / (x^\top x)$ (optionally blended with an item/edge dispersion statistic), producing a single $\lambda$-like score per token per head. Attention logits are then built by comparing the query’s $\lambda_q$ to \textbf{previously stored} $\lambda_k$ values for cached keys:
	 $$a_{ij}=-|\lambda_{q,i}-\lambda_{k,j}|/\mathrm{temp}$$
	after which the causal mask and $V$-weighted sum are reused unchanged. This changes what must be cached: tauformer stores $V$ and the \textbf{scalar} $\lambda_k$ history rather than full $K$, reducing cache size
	\colorbox{yellow!30}{from $2B\,H_{kv}\,T\,D$ floats to $B\,H_{kv}\,T\,D + B\,H_{kv}\,T$ floats}
	(about $\sim50\%$ savings for typical $D$).
	Computing shifts accordingly, instead of spending $O(B\,H\,T\,D)$ on dot-products at each decode step, Tauformer pays $O(B\,H\,\mathrm{nnz}(L))$ to compute the query’s Rayleigh term  with a sparse Laplacian and only $O(B\,H\,T)$ to compare scalars against cached $\lambda_k$, making the scoring step dimension-light while keeping value aggregation $O(B\,H\,T\,D)$ the same (equivalence of outcome compared to \texttt{nanoGPT} are available in the implementation's unit tests).
	
	The key improvement is that the Graph Laplacian is designed to be a sparse matrix, so optimising the code for a sparse representation the $O(D²)$ terms become $O(nnz)$ with $nnz$ being the number of non-zero elements in the matrix. 
	
	Further computing advantages are discussed in \ref{sec:results}.
	
	\subsection{Built-in Memory}
	
	Beside the additional potential computing advantages that are analysed below, the Tauformer and relative \texttt{tauGPT} concept allow the kind of persistent memory described in \ref{sec:generics};  \textbf{delivering a persistent (long-term, unchanging in this initial design) memory within the attention mechanism}. This follows current research for supplementing LLMs with memory sub-systems: for example as seen in \cite{titans:2025} where a neural long-term memory module for 2M+ token contexts works in parallel and provides prefix sub-vectors to the context window; or in \cite{cognitivemem:2025} that categorises memory into sensory, short-term and long-term within agentic workflows; or in \cite{longmem:2024} that uses a "Memory Bank" to store and retrieve long-form text for in-context learning. In addition, the papers that try to supplement LLMs with agentic workflows: for example \cite{amem:2025} that evaluates "Agentic Memory" across multi-hop and temporal reasoning tasks.
	Some of these capabilities can be delivered at a lower-level in the LLM stack. For the purpose of this paper I consider the running attention mechanism the "momentary memory" and the Graph Laplacian against which the distance scores are computed "persistent memory" or Domain-Memory, e.g. a kind of information that is constant to the context as it incorporates some invariants of the frame of reference that is the context in which the generative model generates tokens. This aligns with the scope of my current research for building an "AI Memory Layer" based on topological vector search \cite{tuned_blog:2025}. Tauformer is the first Transformer to natively incorporate some kind of memory capabilities for graph-based data via self-attention. Previous research deals with memory of previous states of the momentary memory like \cite{sukhbaatar:2019} and \cite{bulatov:2022} but not with providing meaningful contextual memory in self-attention. 
	
	Tauformer/tauGPT is the only, at my current knowledge, Transformer architecture that delivers domain knowledge directly in the attention mechanism providing the self-attention itself with some kind of starting context via the knowledge base metadata compressed in \textit{gl}; leveraging the compression made possible by the \textit{arrowspace} library and the \textit{taumode} synthetic index. While technical and conceptual advantages are brought forward by this paper, further research is necessary to ascertain the improved accuracy in the training and token generation on very large context windows and the improved capability in avoiding generation problems sooner than at reasoning time.
	
	\subsection{Code}
	
	The first version of Tauformer/tauGPT \cite{tauformer:2025} has been implemented using the Rust programming language \cite{RustLang:2024} and the Burn Deep-Learning framework \cite{BurnFramework:2025} as they provide fast, structured and type-safe development cycles with mature production ecosystem. Being a concept architecture this code is obviously not meant for production use at the publication of this paper.
	
	\paragraph{No inner-product}

    \texttt{tauGPT}'s most peculiar change is that it swaps dot-product attention for lambda-distance attention, where each token/head vector is compressed into a single scalar $\lambda$ derived from a Laplacian energy, so that attention logits become $-|\Delta\lambda|/T$. The core of compression is in \texttt{lambdas\_from\_heads(...)}, which flattens $[B,H,T,D]$ into $[N,D]$, computes $xL$ via a matmul, then forms $E_{\text{raw}}=(x^\top L x)/(x^\top x+\epsilon)$ and bounds it as $e_raw / (e_raw + tau)$. The scalar $\lambda$ is then broadcast into a full attention matrix using \texttt{taumode\_distance\_logits(...)}, which literally builds logits as $-((lq - lk).abs() / temp)$ after reshaping into $[B,H,Tq,1]$ and $[B,H,1,Tk]$.

	
	\begin{lstlisting}[language=Rust]
	// taumode.rs
	pub fn lambdas_from_heads<B: Backend>(
	x: Tensor<B, 4>,
	lap: Param<Tensor<B, 2>>,
	cfg: &TauModeConfig,
	) -> Tensor<B, 3> {
		let [b, h, t, d] = x.dims();
		
		// Flatten [B,H,T,D] -> [N,D]
		let n = b * h * t;
		let x_nd = x.reshape([n, d]);
		
		// y = x L  -> [N,D]
		let y_nd = x_nd.clone().matmul(lap.val());
		
		// numerator = sum_i x_i * (xL)_i
		let numerator = (x_nd.clone() * y_nd).sum_dim(1); // [N]
		
		// denominator = sum_i x_i^2 + eps
		let denom = x_nd.powf_scalar(2.0).sum_dim(1) + cfg.eps; // [N]
		
		let e_raw = numerator / denom; // [N]
		let e_bounded = e_raw.clone() / (e_raw + cfg.tau); // [N]
		
		e_bounded.reshape([b, h, t])
	}
	
	pub fn taumode_distance_logits<B: Backend>(
	lambda_q: Tensor<B, 3>,
	lambda_k: Tensor<B, 3>,
	cfg: &TauModeConfig,
	) -> Tensor<B, 4> {
		let lq = lambda_q.unsqueeze_dim::<4>(3); // [B,H,Tq,1]
		let lk = lambda_k.unsqueeze_dim::<4>(2); // [B,H,1,Tk]
		let temp = cfg.temperature.max(cfg.eps);
		-((lq - lk).abs() / temp)
	}
	\end{lstlisting}
	
	\paragraph{Sparse Laplacian}
	The other distinctive design choice is that \texttt{TauModeAttention} is built to operate with either a dense test Laplacian or, as intended by design, a sparse Laplacian loaded from a manifold (a \texttt{.parquet} file with the computed Laplacian). The code explicitly switches between those representations at runtime.  In \texttt{TauModeAttention}, you can see the dual storage for the sparse matrix and the dense tensor. All tests are run using the sparse matrix from a pre-trained \texttt{arrowspace} (manifold file and embeddings for the test queries available in \cite{taugpt_kvcache:2025}).
	

	\begin{lstlisting}[language=Rust]
	// tauattention.rs (struct fields + sparse/dense selection)
	pub struct TauModeAttention<B: Backend> {
		// ...
		laplacian_tensor: Option<Param<Tensor<B, 2>>>, // using dense
		laplacian_matrix: Ignored<Option<CsMat<f64>>>, // using sparse
		pub(crate) tau_mode: Ignored<Option<TauMode>>,
		// ...
	}
	
	fn lambdas_from_heads_any(&self, heads: Tensor<B, 4>) -> Tensor<B, 3> {
		let tau_cfg = self.get_tau_config();
		if let Some(lap) = self.laplacian_matrix.0.as_ref() {
			let mode = self.tau_mode.0.unwrap_or(crate::pretraining::parquet::TauMode::Median);
			crate::taumode::lambdas_from_heads_sparse::<B>(heads, lap, mode, tau_cfg.eps)
		} else {
			let lap = self.get_laplacian_tensor().clone();
			crate::taumode::lambdas_from_heads::<B>(heads, lap, &tau_cfg)
		}
	}
    \end{lstlisting}
	
	
	\paragraph{KV-caching layout}
	KV-cache in \texttt{tauGPT} is also different compared to standard GPT: instead of caching $K$ and $V$, each layer caches $(V, lambda_k)$ only, which matches the fact that scoring uses lambda scalars rather than key vectors.  The cache type is defined as \textit{pub type TauCacheLayer<B> = Option<(Tensor<B, 4>, Tensor<B, 3>)>;} and \texttt{tauGPT} wraps a vector of these per layer in \textit{TauKVCache { store: Vec<TauCacheLayer<B>>, position: usize }}.
	
	\begin{lstlisting}[language=Rust]
	// tauattention.rs (KV-cache payload + append logic)
	pub type TauCacheLayer<B> = Option<(Tensor<B, 4>, Tensor<B, 3>)>;
	
	let lambda_k_new = self.lambdas_from_heads_any(k_new); // [B, Hkv, 1]
	
	// Cache management
	let (v_full, lambda_k_full) = match cache_layer.take() {
		Some((v_all, lk_all)) => (
		Tensor::cat(vec![v_all, v_new.clone()], 2),             // time axis
		Tensor::cat(vec![lk_all, lambda_k_new.clone()], 2),     // time axis
		),
		None => (v_new.clone(), lambda_k_new.clone()),
	};
	
	*cache_layer = Some((v_full.clone(), lambda_k_full.clone()));
	let tk = v_full.dims()[2];
	let y = self.scaled_tau_attention_decode(q, lambda_k_full, v_full, 1, tk);
    \end{lstlisting}
	
	 Inside \texttt{TauModeAttention::forward\_decode(...)}, the cache logic appends along the time axis and then runs attention against the full cached history via \texttt{scaled\_tau\_attention\_decode(q, lambda\_k\_full, v\_full, 1, tk)}.
	
	\begin{lstlisting}[language=Rust]
	// taugpt.rs (model-level decode uses cache.position for RoPE step slicing)
	pub fn forward_decode(
	&self,
	last_ids: Tensor<B, 2, Int>, // [B,1]
	cache: &mut TauKVCache<B>,
	use_softcap: bool,
	) -> Tensor<B, 3> {
		let tpos = cache.position;
		let d2 = self.cos.dims()[3];
		
		// Slice RoPE for the current absolute position: [1,1,1,D/2]
		let cos_step = self.cos.clone().slice([0..1, tpos..tpos + 1, 0..1, 0..d2]);
		let sin_step = self.sin.clone().slice([0..1, tpos..tpos + 1, 0..1, 0..d2]);
		
		for (i, block) in self.blocks.iter().enumerate() {
			let layer_cache = &mut cache.store[i];
			x = block.forward_decode(x, (&cos_step, &sin_step), layer_cache);
		}
		
		logits
	}
    \end{lstlisting}
    
    \subsection{Summary of similarities and differences}
    
    Both \texttt{nanoGPT} (\texttt{CausalSelfAttention}) and \texttt{tauGPT} (\texttt{TauModeAttention}) use identical V projection and fetching logic, confirmed by code review and tests.
    
    \paragraph{V projection equivalence}
    
    \textbf{Construction} - both create the same \texttt{c\_v} linear layer with identical initialization:
    
    \begin{lstlisting}[language=Rust]
    cv: LinearConfig::new(n_embd, n_kv_head * head_dim)
    .with_bias(false)
    .with_initializer(KaimingUniform { gain: 0.5, fan_out_only: false })
    \end{lstlisting}

    \textbf{Forward path (prefill)} - both follow the exact same pipeline for V:
    
    \begin{enumerate}
    \item Project: \texttt{self.cv.forward(x).clamp(-5.0, 5.0)}
    \item Reshape: \texttt{.reshape([b, t, self.nkvhead, self.headdim])}
    \item Transpose: \texttt{v.swap\_dims(1, 2)} → \texttt{[B, H\_kv, T, D]}
    \item MQA expand \texttt{(if n\_head != n\_kv\_head)}: both use \texttt{unsqueeze(dim=2).expand(...).reshape(...)} to replicate KV heads
    \item Weighted sum: \texttt{att.matmul(v)} after softmax
    \end{enumerate}
    
    \textbf{Decode path (single-step)} - identical logic for computing new V and caching:
    
    \begin{itemize}
    \item Project step: \texttt{self.cv.forward(x\_step).clamp(-5.0, 5.0).reshape([b, 1, n\_kv\_head, head\_dim]).swap\_dims(1, 2)}
    \item Append to cache: \texttt{Tensor::cat(vec![v\_all, v\_new.clone()], 2)} (time dimension)
    \item MQA expansion and weighted sum identical to forward
    \end{itemize}
    
    \paragraph{What differs}
    The only difference is attention scoring:
    \begin{itemize}
    \item \texttt{nanoGPT}: $Q K^T / \sqrt{d}$ then softmax
    \item \texttt{tauGPT}: lambda-distance logits \texttt{tau\_distance\_logits(lambda\_q, lambda\_k)} then softmax
    \end{itemize}
    
    After that, both apply the same softmax and the same \texttt{att.matmul(v)} to produce output.
    
    \paragraph{Test coverage confirms equivalence}
    Tests verify V is handled identically:
    
    \begin{itemize}
    \item \texttt{test\_forward\_shape\_consistency}: both produce \texttt{[B,T,C]} outputs (V contributes correctly)
    \item \texttt{test\_decode\_single\_step}: both cache V with shape \texttt{[B, H\_kv, T, D]}
    \item \texttt{test\_cache\_accumulation}: V cache grows identically (time dimension increments)
    \item \texttt{test\_mqa\_expansion}: both correctly expand V when \texttt{n\_head != n\_kv\_head}
    \end{itemize}
    
    The difference in final outputs (\texttt{test\_tau\_vs\_causal\_outputs\_differ\_expected}) is purely due to scoring

	\section{Results}
	\label{sec:results}
	
	Multiple tests on both attention mechanisms have been run and results have been collected using \cite{taugpt_kvcache:2025}.
	
	Here some data and diagrams from the tests results. All this data was collected running tests on CPU hardware because the objective of this concept architecture's test results are preliminary and to be taken as a reference and a baseline, to demonstrate that Tauformer/tauGPT can work as any other Transformer but with delivering of domain-specific (topological) metadata directly in the attention mechanism and with some performance gains for larger windows and higher dimensional embeddings.  
	
	Considering generated tokens, the results contains 500000 decoded-token latency rows spanning both engines ("nano", "tau"), both modes (\texttt{kv\_cache}, \texttt{no\_cache}), and \texttt{gen\_tokens} (context window) $\in {128, 256, 512, 1024, 2048}$.  Per-token comparisons were computed by matching "nano" and "tau" tokens on \texttt{prompt\_id}, \texttt{mode}, \texttt{gen\_tokens}, \texttt{token\_index} and then taking the ratio,

	The full normalised token-latency results confirms that \texttt{tauGPT} is consistently faster in "kv-cache" mode (a proxy for inference decoding), and consistently slower in "no cache" (training/uncached proxy) for small context windows (up to 2048 tokens); with the gap widening in \texttt{tauGPT}'s favour as generation length grows in "kv-cache" while the gap in "no cache" mode gets smaller (see \ref{fig:window_grow}). The difference in "no cache" between the two gets smaller as the lengths grow so there is room to make the hypothesis that at some threshold where context window and embeddings lengths are large enough \texttt{tauGPT} will be at least as fast as \texttt{nanoGPT} in training while being increasingly faster at inference time. Also to consider that \texttt{nanoGPT} is a relatively much more mature implementation. This could confirm that Tauformer could be viable also for training on multimodal data where multiple modalities are concatenated in the same embedding space while keeps being faster and more efficient (-50\% memory for KV-cache) at inference. The assumption that gains are provided as the context window and the embeddings dimensions grows is for now supported only by regressions run on ~0.5M decoded-token rows across hundreds of runs of testing with different lengths. In particular an idea about at which threshold \textit{tauGPT} could reach equivalent performance also in training is hinted in \ref{fig:window_grow}, \ref{fig:prompt_grow}, \ref{fig:nembd_grow}.
	
	\colorbox{yellow!30}{See table \ref{tab:benchmark-summary} for summary statistics} (repository \cite{taugpt_kvcache:2025} may contain more up-to-date data in the \texttt{report/} directory).
	
	TTFT (Time To First Token) is measured differently depending on the caching mode:
	\begin{itemize}
	\item "no\_cache" mode: TTFT includes the full forward pass through the entire prompt context, since there's no cached state. The timer starts when processing begins and stops when the first token is generated; \texttt{end\_to\_end\_ms = ttft\_ms + decode\_total\_ms}.
	\item "kv\_cache" mode: TTFT measures only the decode step for producing the first new token, after the cache has already been built from the prompt. The prefill/priming time is tracked separately in the \texttt{prefill\_ms} and \texttt{prime\_ms columns}; \texttt{end\_to\_end\_ms = prefill\_ms + ttft\_ms + decode\_total\_ms}.
	\end{itemize}
	
	\begin{table*}[htbp]
		\centering
		\caption{Benchmark Summary Statistics for low dimensional vectors}
		\label{tab:benchmark-summary}
		\begin{tabular}{@{}llrrr@{}}
			\toprule
			\multicolumn{5}{c}{\textbf{Throughput (tokens/sec) -- Higher is Better}} \\
			\midrule
			\textbf{Engine} & \textbf{Mode} & \textbf{Mean} & \textbf{Min} & \textbf{Max} \\
			\midrule
			nano & kv\_cache & 37.47 & 13.91 & 59.96 \\
			nano & no\_cache & 4.55 & 0.89 & 9.71 \\
			tau & kv\_cache & 49.39 & 15.04 & 71.21 \\
			tau & no\_cache & 3.43 & 0.81 & 6.91 \\
			\midrule
			\multicolumn{5}{c}{} \\
			\multicolumn{5}{c}{\textbf{Per-Token Latency (ms) -- Lower is Better}} \\
			\midrule
			\textbf{Engine} & \textbf{Mode} & \textbf{Mean} & \textbf{p50} & \textbf{p95} \\
			\midrule
			nano & kv\_cache & 36.25 & 36.26 & 41.35 \\
			nano & no\_cache & 327.58 & 295.88 & 550.58 \\
			tau & kv\_cache & 30.44 & 30.42 & 33.13 \\
			tau & no\_cache & 390.04 & 359.90 & 636.52 \\
			\bottomrule
		\end{tabular}
		\quad
	\begin{tabular}{@{}lr@{}}
		\toprule
		\multicolumn{2}{c}{\textbf{Speedup: Tau/Nano}} \\
		\multicolumn{2}{c}{($>1.0$ = Tau faster)} \\
		\midrule
		\textbf{Metric} & \textbf{Factor} \\
		\midrule
		Tokens/sec & 1.068$\times$ \\
		TTFT & 0.997$\times$ \\
		p50 latency & 1.066$\times$ \\
		p95 latency & 1.104$\times$ \\
		Decode time & 1.067$\times$ \\
		Total time & 1.028$\times$ \\
		\bottomrule
	\end{tabular}

	\end{table*}
	
	\subsection{Inference proxy (kv\_cache) results}
	
	On a matched per-token basis, nano/tau speedup medians (p50) increase with \texttt{gen\_tokens}: 1.09× (128), 1.12× (256), 1.17× (512), and 1.22× (1024), meaning tau's per-token latency advantage strengthens as sequences get longer. Tail behaviour also improves in tau's favour as sequences grow: the p95 of nano/tau per-token speedup rises from 1.35× (128) to 1.52× (1024), suggesting tau reduces worse-case token stalls relative to nano at longer generations in "kv-cache" mode.  Looking specifically at the last 25\% of tokens (75–100\% bins), the median-latency ratio nano/tau reaches ~1.33× at \texttt{gen\_tokens=1024}, indicating the end of long generations is where tau helps most under \texttt{kv\_cache}.
	
	\subsection{Training/uncached proxy (no\_cache) results}
	
	For "no cache", matched per-token speedups are $\le1$ across all lengths (tau slower): median nano/tau speedup is 0.75× (128), 0.75× (256), 0.76× (512), and 0.80× (1024).  The late-token-bin summary also stays $\le1$ for "no cache" (e.g., last-25\% median ratio is $\sim0.86×$ at 1024), but the ratio moves upward with length, which suggests tauGPT's relative penalty shrinks at longer sequences even without caching.  In the "no cache" distribution, even the p99 nano/tau speedup stays below 1 ($\simeq0.95–0.97$), meaning tauGPT is slower almost everywhere token-by-token in this mode rather than losing only in a small tail, the situation improves when the window and embeddings lengths grow. Considering the results of the regressions for larger lengths (context window, embeddings dimensions and prompt length $>10^5$) there is still a lot of margin to harvest in training also considering that this implementation is at concept stage without any optimisation. Results at 2048 are already better promising even if the number of runs is limited at the moment.
	
	\subsection{Matched per-token speedup distribution (nano/tau)}
	
	This is the test matrix and the relative results:
	
	\begin{table}[H] % [H] forces the table to stay exactly here
		\centering
		\begin{tabular}{l r r r r r}
			\hline
			\textbf{mode} & \textbf{gen\_tokens} & \textbf{matched token pairs} & \textbf{speedup p50} & \textbf{speedup p95} & \textbf{speedup p99} \\
			\hline
			kv\_cache & 128  & 96{,}768  & 1.093 & 1.346 & 1.437 \\
			kv\_cache & 256  & 165{,}888 & 1.124 & 1.395 & 1.494 \\
			kv\_cache & 512  & 62{,}208  & 1.166 & 1.443 & 1.557 \\
			kv\_cache & 1024 & 359{,}424 & 1.223 & 1.523 & 1.658 \\
			no\_cache & 128  & 96{,}768  & 0.747 & 0.889 & 0.965 \\
			no\_cache & 256  & 165{,}888 & 0.749 & 0.898 & 0.968 \\
			no\_cache & 512  & 62{,}208  & 0.762 & 0.901 & 0.954 \\
			no\_cache & 1024 & 359{,}424 & 0.805 & 0.924 & 0.965 \\
			\hline
		\end{tabular}
		\caption{Matched per-token speedup distribution (nano/tau) across cache modes and sequence lengths (higher numbers mean tau is better).}
		\label{tab:token-speedup-distribution}
	\end{table}
	
	These diagrams compare the gains as the context window, the prompt length and embeddings dimensions grow up beyond $\sim10^7$, this is only indicative as for \texttt{nanoGPT} is running on CPU and \texttt{tauGPT} is unoptimised and a concept model but I think it can provides hints to frame the orders of magnitude at which Tauformer can become performant also in \texttt{no\_cache} mode or in other scenarios:
	
	\begin{figure}[H]
		\centering
		\includegraphics[width=0.8\textwidth]{tau_speedup_vs_gentokens_bothmodes.png}
		\caption{Regressions for tau/nano speedups as the context window grows}
		\label{fig:window_grow}
	\end{figure}
	
	\begin{figure}[H]
		\centering
		\includegraphics[width=0.8\textwidth]{tau_speedup_vs_promptlen_bothmodes.png}
		\caption{Regressions for tau/nano speedups as the prompt length grows}
		\label{fig:prompt_grow}
	\end{figure}
	
	\begin{figure}[H]
		\centering
		\includegraphics[width=0.8\textwidth]{tau_speedup_vs_nembd_bothmodes.png}
		\caption{Regressions for tau/nano speedups as embeddings' dimensions grows}
		\label{fig:nembd_grow}
	\end{figure}

	\section{Conclusion: A Beginning}
	
	Potential for larger improvements are tested in larger windows generation tests (256, 512, 1024, ..., $10^6$ tokens) and larger embeddings dimensions. The driving hypothesis is that when caching becomes too expensive memory-wise (like for very long windows with high-dimensional embeddings), \texttt{tauGPT} can provide a better solution both in terms of contextual accuracy and performance. Current research is exploring at the order of magnitude of $10^4$ window length (average for LLMs is 6000 to 14000) with values for embeddings dimensions in the same order of magnitude. The idea is to test to have a prompt length and embeddings length that are of the same order of magnitude of number of non-zero values of the sparse representation in the pre-trained manifold, so that the model can adapt to the number of dimensions of the manifold processing an element in a single query ($Q$). The prompt length should be obviously modifiable at will.
	
	\textit{The normalised token latencies indicate tau's current value proposition is "decode-heavy inference on long windows and high dimensions" with promising performance gains both in a "no cache" and "kv-cache" mode for large scales}.  If training work is well proxied by \texttt{no\_cache}, \texttt{tauGPT} currently underperforms across essentially the entire per-token distribution on lower scales; so demonstrating training wins require architectural/implementation optimisations, or for the model to be tested on very long context windows and very high embeddings dimensions (as showed in the regression diagrams). The next experiments that would most strengthen Tauformer position should aim to achieve the same results beyond lengths of $10^7$ for the $H$ $D$ and  $T$.
	
	Current results demonstrate the promising characteristics of a Topological Transformer architecture, Tauformer; designed to deliver
	\begin{itemize}
	\item \textit{(i)} domain-specific metadata at the level of attention computation and
	\item \textit{(ii)} make self-attention more convenient at very large "lengths" (sequences, heads, embeddings dimensions) leveraging topological search (using \texttt{arrowspace}'s \texttt{taumode} vector compression).
	\end{itemize}
	To synthesise in few words, Tauformer leverages \texttt{taumode} as attention score, allowing to avoid usage of inner-products in the attention computation, potentially making possible very large context windows and to run attention on higher dimensional (multimodal) embeddings with accessible costs and infrastructural requirements on multi nodes hardware.
	
	In the current implementation, \texttt{TauModeAttention} still forms attention logits/probabilities over the whole key-time axis and then does the usual weighted sum \texttt{att.matmul(v)}. That means decode-time work and intermediate memory scale like $O(T)$ per generated token (and $O(H \cdot T)$ across heads), which becomes prohibitive at $T = 10^7$ unless of a redesign of the attention computation to be streaming/chunked/sparse that is one of the principles stated (to use sparse representation).
	
	To benefit at $10^9$ context, Tauformer needs a variant where it does not cache raw per-token $V$ (even compressed) and then multiply by a length-$10^9$ attention vector.	Instead, cache per-centroid aggregates (like $S\_c, U\_c$) and compute attention over $K$ centroids (or over a small retrieved subset), so both memory and compute scale with $K$, not $T$. This idea brings forward the possibility of compressing the historical $V$ using \texttt{taumode} and \texttt{gl} for the latent space as well: for example, use the same Laplacian idea already applied to get $\lambda$ scalars (via Laplacian-based mapping helper) but apply it as a projection of $V$ onto a low-dimensional basis (e.g., first $r$ spectral components), store only coefficients, and reconstruct an approximate $V$ when needed. This is analogous to how \texttt{arrowspace} does dimensionality reduction before computing the Laplacian but in the latent/value manifold rather than in token embedding space. This is another potential example of the hypothetical members of a Topological Transfomers' class of models built for very large lengths.
	
	\section{Acknowledgments}
	The author is an independent researcher who self-funded this work. All works available at \href{https://www.tuned.org.uk}{his research page}. Thanks to my supporting network of peers and to all the backers and sponsors to my research effort.
	
	\newpage
	
	\bibliographystyle{plain}
	\begin{thebibliography}{15}
		\bibitem{ArrowSpace:2025} Lorenzo Moriondo, \newblock \textit{ArrowSpace: introducing Spectral Indexing for vector search}, 2025. \newblock \url{https://github.com/tuned-org-uk/arrowspace-rs}
		DOI: 10.21105/joss.09002

		\bibitem{nanoGPT:2022} Andrej Karpathy
        \newblock {\em nanoGPT: The simplest, fastest repository for training open-source GPT models}. \url{https://github.com/karpathy/nanoGPT}
        \newblock
		
		\bibitem{vaswani_attention:2017} Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and Kaiser, Lukasz and Polosukhin, Illia \newblock \textit{Attention is All You Need}, Advances in Neural Information Processing Systems (NeurIPS), volume 30, 2017 \newblock

		\bibitem{Rayleigh:1877} J.~W.~Strutt (Lord Rayleigh), \textit{The Theory of Sound}, Vol. 1, London: Macmillan and Co., 1877.

		\bibitem{wang-etal-tsdae:2021} Kexin Wang, Nils Reimers, and Iryna Gurevych. 2021. 
        \newblock TSDAE: Using Transformer-based Sequential Denoising Auto-Encoder for Unsupervised Sentence Embedding Learning. 
        \newblock In \emph{Findings of the Association for Computational Linguistics: EMNLP 2021}, pages 671--688, Punta Cana, Dominican Republic. Association for Computational Linguistics.

        \bibitem{lefebvre_sepal:2025} Félix Lefebvre and Gaël Varoquaux. 2025.
        \newblock SEPAL: Scalable Embedding Propagation ALgorithm for large knowledge graphs.
        \newblock In \emph{Proceedings of the 13th International Conference on Learning Representations (ICLR 2025)}, Singapore.
		
		\bibitem{tauformer:2025} Lorenzo Moriondo (@Mec-iS)
		\newblock { Tauformer }. \url{		https://github.com/tuned-org-uk/taugpt-kvcache-bench}
		\newblock
		
		\bibitem{taugpt_kvcache:2025} Lorenzo Moriondo (@Mec-iS)
		\newblock { KV-cache bench for Tauformer}. \url{		https://github.com/tuned-org-uk/taugpt-kvcache-bench}
		\newblock
		
		\bibitem{titans:2025}
		Ali Behrouz, Peilin Zhong, and Vahab Mirrokni.
		\newblock Titans: Learning to Memorize at Test Time.
		\newblock {\em arXiv preprint arXiv:2501.00663}, 2025.
		
		\bibitem{cognitivemem:2025}
		Lianlei Shan, Shixian Luo, Zezhou Zhu, Yu Yuan, Yong Wu
		\newblock Cognitive Memory in Large Language Models.
		\newblock {\em arXiv preprint arXiv:2504.02441}, 2025.
		
		\bibitem{longmem:2024}
		Weizhi Wang, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, Furu Wei
		\newblock Augmenting Language Models with Long-Term Memory.
		\newblock {\em arXiv preprint arXiv:2306.07174}, 2024 (Updated).
		
		\bibitem{amem:2025}
		Wujiang Xu, Zujie Liang, Kai Mei, Hang Gao, Juntao Tan, Yongfeng Zhang
		\newblock A-MEM: Agentic Memory for LLM Agents.
		\newblock {\em arXiv preprint arXiv:2502.12110}, 2025.
		
		\bibitem{tuned_blog:2025}
		Blog posts at tuned.org.uk.
		\newblock Tuned Blog.
		\newblock \url{www.tuned.org.uk/blog}, 2025.
		\newblock Accessed: December 29, 2025.
	
		\bibitem{sukhbaatar:2019}
		S.~Sukhbaatar, E.~Grave, G.~Lample, H.~Jegou, and A.~Joulin,
		\newblock Augmenting Self-attention with Persistent Memory
		\newblock \emph{arXiv preprint arXiv:1907.01470}, 2019.
		doi:10.48550/arXiv.1907.01470.
		
		\bibitem{bulatov:2022}
		A.~Bulatov, Y.~Kuratov, and M.~S.~Burtsev,
		\newblock Recurrent Memory Transformer
		\newblock \emph{arXiv preprint arXiv:2207.06881}, 2022.
		doi:10.48550/arXiv.2207.06881.
		
		\bibitem{RustLang:2024}
		The Rust Project Developers, \textit{The Rust Programming Language}, 2024. [Online]. Available: \url{www.rust-lang.org}
		
		\bibitem{BurnFramework:2025}
		Tracel-AI, \textit{Burn: A Deep Learning Framework Designed from Engineers' Perspectives}, 2025. [Online]. Available: \url{https://github.com/tracel-ai/burn}
		
	\end{thebibliography}
	
	
\end{document}