Papers
Topics
Authors
Recent
Search
2000 character limit reached

Kan Extension Transformers (KETs)

Updated 5 July 2026
  • 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

F:CE,K:CD,F : \mathcal{C} \to \mathcal{E}, \qquad K : \mathcal{C} \to \mathcal{D},

with the left Kan extension of FF along KK, denoted LanKF\operatorname{Lan}_K F, defined as a universal extension

LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}

equipped with a natural transformation

η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K

such that for any

G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,

there exists a unique natural transformation

α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G

with

γ=(αK)η.\gamma = (\alpha K)\circ \eta.

Dually, the right Kan extension RanKF\operatorname{Ran}_K F is characterized by maps

FF0

The paper further states the pointwise colimit and limit formulas

FF1

and the enriched coend and end expressions

FF2

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

FF3

and source-side features are assigned by

FF4

This yields the common weighted aggregation form

FF5

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” FF6 encode compatibility or incidence between source object FF7 and target token FF8, 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: FF9 Each source object is a token KK0, and the update at target token KK1 is

KK2

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 KK3 parameterizing weights KK4 against keys KK5, and values KK6 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

KK7

The resulting KK8 is described as a learned geometric graph with fuzzy KK9-NN weights LanKF\operatorname{Lan}_K F0 induced by topological proximity rather than sequence position alone. Its update is

LanKF\operatorname{Lan}_K F1

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 LanKF\operatorname{Lan}_K F2-simplices, to edges, viewed as LanKF\operatorname{Lan}_K F3-simplices, and optionally to higher-order simplices such as faces or motifs. The generic simplicial update is

LanKF\operatorname{Lan}_K F4

Two concrete realizations are specified. The first is quadratic KET: LanKF\operatorname{Lan}_K F5 which is described as an attention kernel over simplices rather than over tokens and is typically LanKF\operatorname{Lan}_K F6 in sequence length. The second is incidence-restricted KET: LanKF\operatorname{Lan}_K F7 with the edge-only causal example

LanKF\operatorname{Lan}_K F8

This reduces complexity to LanKF\operatorname{Lan}_K F9 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 LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}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

LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}1

with

LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}2

Denoising completion is formulated as

LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}3

together with the corruption schedule

LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}4

which spans corruption rates from LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}5 to LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}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

LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}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

LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}8

and equivalently in the appendix as

LanKF:DE\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}9

Here η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K0 is the embedding matrix and η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K1 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 η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K2. 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 η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K3 softmax attention. TopoCoend performs token-level aggregation over learned topological neighborhoods, with complexity depending on η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K4-NN sparsity and weights derived from latent η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K5 proximity. Quadratic KET uses a global attention-like kernel over simplices and is typically η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K6, 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 η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K7; 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 η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K8, a simplex set η:F(LanKF)K\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K9, a value base G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,0, and an optional causal mask; it constructs simplex values G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,1, simplex keys G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,2, computes token queries G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,3, optionally restricts the simplex set to G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,4 under a causal mask, forms

G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,5

aggregates

G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,6

and updates

G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,7

The Incidence-Restricted Kan Block (Edge-Only) computes

G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,8

for G:DE,γ:FGK,G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,9, then initializes α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G0, adds α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G1 when α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G2, optionally adds α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G3 when the regime is noncausal and α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G4, and updates

α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G5

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 α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G6 (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 α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G7, weight decay α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G8, batch size α:LanKFG\alpha : \operatorname{Lan}_K F \Rightarrow G9, context length γ=(αK)η.\gamma = (\alpha K)\circ \eta.0, γ=(αK)η.\gamma = (\alpha K)\circ \eta.1 layers, hidden dimension γ=(αK)η.\gamma = (\alpha K)\circ \eta.2, γ=(αK)η.\gamma = (\alpha K)\circ \eta.3 attention heads, γ=(αK)η.\gamma = (\alpha K)\circ \eta.4 training iterations, reporting every γ=(αK)η.\gamma = (\alpha K)\circ \eta.5 steps, evaluation every γ=(αK)η.\gamma = (\alpha K)\circ \eta.6 steps, self-conditioned temperature γ=(αK)η.\gamma = (\alpha K)\circ \eta.7, and for TopoCoend γ=(αK)η.\gamma = (\alpha K)\circ \eta.8, γ=(αK)η.\gamma = (\alpha K)\circ \eta.9. Structured completion runs use matched Transformer and incidence-KET backbones, AdamW, learning rate RanKF\operatorname{Ran}_K F0, weight decay RanKF\operatorname{Ran}_K F1, context length RanKF\operatorname{Ran}_K F2, hidden dimension RanKF\operatorname{Ran}_K F3, RanKF\operatorname{Ran}_K F4 and RanKF\operatorname{Ran}_K F5 layers, RanKF\operatorname{Ran}_K F6 attention heads, block size RanKF\operatorname{Ran}_K F7, RanKF\operatorname{Ran}_K F8 training steps, evaluation every RanKF\operatorname{Ran}_K F9 steps, denoising steps FF00, and gradients clipped to norm FF01 (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 FF02, WT2 FF03, and WT103 FF04. For KET-Inc-C to KET-Inc-PD they are PTB FF05, WT2 FF06, and WT103 FF07. GT-PD attains PTB FF08, WT2 FF09, and WT103 FF10 (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 FF11, confirming that validity depends on which values flow through the graph (Mahadevan, 26 May 2026).

For structured language modeling with FF12, context FF13, and FF14, direct block prediction is reported to be much harder than denoising completion. Representative values include PTB with TF-Block-4 FF15 at FF16 versus KET-Block-4 FF17 at FF18, while TF-Denoise-4 FF19 gives FF20 and KET-Denoise-4 FF21 gives FF22. On WT2 with FF23, TF-Block-4 gives FF24, KET-Block-4 FF25, TF-Denoise-4 FF26, and KET-Denoise-4 FF27. On WT103 with FF28, TF-Block-4 gives FF29, KET-Block-4 FF30, TF-Denoise-4 FF31, and KET-Denoise-4 FF32 (Mahadevan, 26 May 2026). The stated interpretation is that direct block prediction into FF33 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.

Definition Search Book Streamline Icon: https://streamlinehq.com
References (1)

Topic to Video (Beta)

No one has generated a video about this topic yet.

Whiteboard

No one has generated a whiteboard explanation for this topic yet.

Follow Topic

Get notified by email when new papers are published related to Kan Extension Transformers (KETs).