Kan Extension Transformers (KETs)
- KETs are a categorical formulation where Transformer layers aggregate features via weighted structured extension operators, unifying self-attention, geometric, and simplicial methods.
- They extend standard self-attention by incorporating higher-order simplicial structures and learned topological neighborhoods to capture richer contextual relationships.
- Empirical evaluations show that applying predict-detach regimes in KET architectures significantly improves performance by controlling information flow and preventing gradient leakage.
Kan Extension Transformers (KETs) are a categorical formulation of Transformer layers in which contextualization is expressed as a weighted structured extension operator. In this framework, a layer aggregates values carried by source objects defined over a neighborhood system and extends them to target token positions; standard self-attention appears as the singleton-neighborhood case, Geometric Transformer style incidence mixing appears as a sparse edge-restricted case, and KET constitutes the higher-order simplicial case (Mahadevan, 26 May 2026). The same formal lens is used to relate aggregation to diffusion-style completion and to characterize predict-detach self-conditioning as a causally valid mechanism for exposing noncausal structure without transporting teacher-forced future hidden states (Mahadevan, 26 May 2026).
1. Categorical definition and formal apparatus
The core categorical setup introduces functors
with the left Kan extension of along , denoted , defined as a universal extension
equipped with a natural transformation
such that for any
there exists a unique natural transformation
with
Dually, the right Kan extension is characterized by maps
0
The paper further states the pointwise colimit and limit formulas
1
and the enriched coend and end expressions
2
The interpretive summary given in the paper is concise: left Kan corresponds to aggregation from local evidence, whereas right Kan corresponds to completion under constraints (Mahadevan, 26 May 2026).
Within the Transformer setting, the target token positions are written as
3
and source-side features are assigned by
4
This yields the common weighted aggregation form
5
The terminology is explicit. “Carriers” are the values transported from source objects, including token states, edge embeddings, and higher-order simplex values. “Neighborhoods” are the objects over which aggregation occurs, including tokens, learned topological neighbors, and simplices. “Weights” 6 encode compatibility or incidence between source object 7 and target token 8, including attention kernels, geometric proximity, and simplicial incidence (Mahadevan, 26 May 2026).
This formulation is not merely a restatement of attention in categorical notation. It is intended as a common language in which several Transformer-like mechanisms become instances of the same extension schema. A plausible implication is that architectural variation can be analyzed along two largely separable axes: the source neighborhood family and the information regime carried through that family.
2. Attention, geometric incidence mixing, and simplicial generalization
Standard self-attention is recovered by restricting the source neighborhood system to tokens: 9 Each source object is a token 0, and the update at target token 1 is
2
Categorically, this is presented as a weighted left-Kan-style extension in which source and target objects are token positions; the familiar query-key-value mechanism supplies the weights and carriers, with queries 3 parameterizing weights 4 against keys 5, and values 6 carrying the features associated with source objects (Mahadevan, 26 May 2026).
The geometric case, called TopoCoend, preserves tokens as source objects but changes the neighborhood structure from positional adjacency to learned geometric adjacency in a latent space
7
The resulting 8 is described as a learned geometric graph with fuzzy 9-NN weights 0 induced by topological proximity rather than sequence position alone. Its update is
1
This remains a coend-style weighted extension, but with learned neighborhoods rather than a fixed sequential token-to-token pattern (Mahadevan, 26 May 2026).
KET changes the source objects themselves by lifting from tokens, viewed as 2-simplices, to edges, viewed as 3-simplices, and optionally to higher-order simplices such as faces or motifs. The generic simplicial update is
4
Two concrete realizations are specified. The first is quadratic KET: 5 which is described as an attention kernel over simplices rather than over tokens and is typically 6 in sequence length. The second is incidence-restricted KET: 7 with the edge-only causal example
8
This reduces complexity to 9 and is used to clarify the connection to Geometric Transformer, which is identified as an incidence-restricted special case of KET (Mahadevan, 26 May 2026).
A significant restriction is also stated explicitly. Although the framework is simplicial and topological in its indexing, the KET equations do not introduce explicit boundary or coboundary operators 0, Hodge Laplacians, or weighted cochain complexes. The simplicial structure enters through neighborhood indexing and incidence-restricted transport rather than through explicit Laplacian or chain-complex formulas (Mahadevan, 26 May 2026). This point addresses a likely misconception: KET is not presented as a Hodge-theoretic or message-passing Laplacian model, even though it invokes simplices and incidence.
3. Aggregation, completion, and the diffusion-style bridge
The paper connects left Kan aggregation to right Kan completion by introducing structured block prediction and denoising. Direct block prediction is formulated as
1
with
2
Denoising completion is formulated as
3
together with the corruption schedule
4
which spans corruption rates from 5 to 6 (Mahadevan, 26 May 2026).
The conceptual claim is that left Kan constructions, expressed through coends, describe aggregation from local evidence, whereas right Kan constructions, expressed through ends or limits, describe completion subject to multiple local constraints. Denoising block prediction uses partial future objects as constraints and asks the model to fill in missing structure; the paper describes this using a “horn filling” analogy from simplicial homotopy (Mahadevan, 26 May 2026).
The formulation is deliberately limited. The text states that it does not present explicit diffusion updates such as
7
Instead, denoising is positioned as a right-Kan-style completion regime rather than as a discretization of a Laplacian flow or a particular stochastic differential process (Mahadevan, 26 May 2026). This distinction matters because it prevents over-identifying the framework with standard diffusion implementations. The bridge is structural and conceptual: completion under constraints resembles denoising, but the paper does not reduce KET to a specific diffusion algorithm.
The structured-completion perspective broadens the scope of the framework beyond autoregressive language modeling. The paper states that it generalizes to blockwise infilling and broader tasks where partial target structure is available (Mahadevan, 26 May 2026). This suggests a research program in which the same extension operator is reused across causal prediction, block prediction, and denoising completion, with the primary change occurring in the information supplied to the operator.
4. Predict-detach self-conditioning and causal validity
A central mechanism is predict-detach self-conditioning, in which noncausal neighborhoods are made compatible with causal validity by changing what is transported through those neighborhoods. The detached predictive carrier is defined as
8
and equivalently in the appendix as
9
Here 0 is the embedding matrix and 1 is a temperature parameter (Mahadevan, 26 May 2026).
The rationale for detach is given in two explicit points. First, there are “No gold-future values”: the carrier is computed from a prefix-valid predictive state rather than from a hidden state that already encodes the true future under teacher forcing. Second, there is “No leakage gradient”: detaching blocks the auxiliary branch from becoming a backdoor through which the model could cheaply encode targets into the carriers that will later be consumed by noncausal aggregation (Mahadevan, 26 May 2026).
Under this regime, the same extension operator—whether attention, topological aggregation, or simplicial aggregation—can reuse noncausal neighborhoods without violating causal validity because the transported values are endogenous predictions frozen by detach rather than teacher-forced future states (Mahadevan, 26 May 2026). The point of emphasis is that validity depends on which values flow through the graph, not merely on whether the graph itself includes noncausal edges.
The paper includes a leakage diagnostic. Targets are shuffled within a batch; if noncausal neighborhoods transport teacher-forced hidden states, perplexity collapses toward 2. By contrast, predict-detach regimes remain high under shuffling, which is taken to indicate causal validity (Mahadevan, 26 May 2026). This diagnostic supports the paper’s stronger claim that the principal issue is information leakage through carriers, not noncausal connectivity as such.
A common misunderstanding is therefore directly addressed by the reported experiments: noncausal structure is not inherently invalid. According to the paper, it becomes valid when paired with detached predictive carriers and invalid when paired with teacher-forced future states (Mahadevan, 26 May 2026).
5. Computational regimes, implementations, and complexity
The framework distinguishes several computational regimes by neighborhood family and aggregation rule. Attention over tokens is standard 3 softmax attention. TopoCoend performs token-level aggregation over learned topological neighborhoods, with complexity depending on 4-NN sparsity and weights derived from latent 5 proximity. Quadratic KET uses a global attention-like kernel over simplices and is typically 6, with greater expressivity arising from higher-order source objects. Incidence-restricted KET performs local aggregation from incident simplices such as edges and faces and is 7; it is described as an efficient sparse approximation. GT is explicitly categorized as an incidence-restricted special case of KET (Mahadevan, 26 May 2026).
The paper remarks that quadratic KET benefits disproportionately from modern GPU acceleration, which narrows its runtime gap relative to incidence-based variants, while incidence-restricted KET and GT are lighter-weight linear-time mixers (Mahadevan, 26 May 2026). This is presented as a practical trade-off rather than a formal complexity-theoretic novelty.
Two implementation blocks are specified verbatim. The Quadratic Kan Extension Block takes token states 8, a simplex set 9, a value base 0, and an optional causal mask; it constructs simplex values 1, simplex keys 2, computes token queries 3, optionally restricts the simplex set to 4 under a causal mask, forms
5
aggregates
6
and updates
7
The Incidence-Restricted Kan Block (Edge-Only) computes
8
for 9, then initializes 0, adds 1 when 2, optionally adds 3 when the regime is noncausal and 4, and updates
5
The paper also states that initialization follows standard Transformer practice, while stability relies on detach to block gradients through the auxiliary branch and on gradient clipping with norm 6 (Mahadevan, 26 May 2026).
This operational detail reinforces the paper’s general distinction between structural design and information regime. The same block structure can behave either causally or noncausally depending on the carriers and gradient pathways.
6. Empirical evaluation, interpretation, and limitations
The empirical study evaluates 12 different Transformer implementations across strict-causal and predict-detach regimes on Penn Treebank, WikiText-2, and WikiText-103 (Mahadevan, 26 May 2026). For the main causal language-model comparisons, the reported configuration uses AdamW, learning rate 7, weight decay 8, batch size 9, context length 0, 1 layers, hidden dimension 2, 3 attention heads, 4 training iterations, reporting every 5 steps, evaluation every 6 steps, self-conditioned temperature 7, and for TopoCoend 8, 9. Structured completion runs use matched Transformer and incidence-KET backbones, AdamW, learning rate 0, weight decay 1, context length 2, hidden dimension 3, 4 and 5 layers, 6 attention heads, block size 7, 8 training steps, evaluation every 9 steps, denoising steps 00, and gradients clipped to norm 01 (Mahadevan, 26 May 2026).
The strict-causal test perplexities are reported as follows:
| Model | PTB | WT2 | WT103 |
|---|---|---|---|
| Transformer | 124.47 | 163.92 | 232.52 |
| GT-Causal | 127.17 | 157.74 | 215.69 |
| KET-Quad-C | 133.37 | 156.42 | 210.30 |
| KET-Inc-C | 137.19 | 161.12 | 213.76 |
The interpretation given in the paper is that on PTB, the plain Transformer is strongest among strict-causal models, whereas on WT2 and WT103, quadratic KET is the strongest among the tabulated causal architectures (Mahadevan, 26 May 2026).
The predict-detach ablation is more pronounced. Test perplexity changes from KET-Quad-C to KET-Quad-PD are PTB 02, WT2 03, and WT103 04. For KET-Inc-C to KET-Inc-PD they are PTB 05, WT2 06, and WT103 07. GT-PD attains PTB 08, WT2 09, and WT103 10 (Mahadevan, 26 May 2026). The paper’s interpretation is that the largest empirical gains come from the predict-detach regime itself, while neighborhood family determines how those gains are expressed. It further reports that gold noncausal leakage diagnostics, in which teacher-forced future states are passed through noncausal neighborhoods, produce collapse toward perplexity 11, confirming that validity depends on which values flow through the graph (Mahadevan, 26 May 2026).
For structured language modeling with 12, context 13, and 14, direct block prediction is reported to be much harder than denoising completion. Representative values include PTB with TF-Block-4 15 at 16 versus KET-Block-4 17 at 18, while TF-Denoise-4 19 gives 20 and KET-Denoise-4 21 gives 22. On WT2 with 23, TF-Block-4 gives 24, KET-Block-4 25, TF-Denoise-4 26, and KET-Denoise-4 27. On WT103 with 28, TF-Block-4 gives 29, KET-Block-4 30, TF-Denoise-4 31, and KET-Denoise-4 32 (Mahadevan, 26 May 2026). The stated interpretation is that direct block prediction into 33 is much harder than denoising completion, again highlighting regime effects over pure architecture choice.
The paper’s practical guidance follows directly from these results. In strict-causal language modeling, neighborhood design matters: quadratic KET is strongest among the compared causal architectures on WT2 and WT103, and incidence-restricted KET remains competitive with lower cost. Across all datasets, however, the larger gains come from the predict-detach regime rather than from changing the neighborhood family alone (Mahadevan, 26 May 2026).
The limitations are also explicit. The study is modest in scale, centered on perplexity, and uses single-seed runs. Future work is said to require scaling neighborhood design and information regime separately, exploring richer simplicial families, adaptive simplex selection, longer-context sparse approximations, and hybrid systems combining KET with TopoCoend (Mahadevan, 26 May 2026). This suggests that the current contribution is primarily a unifying formalism with targeted empirical support rather than a definitive large-scale benchmark study.