Kan Extension Transformers: A Categorical Unification of Attention, Diffusion, and Predict-Detach Self-Conditioning

Abstract: We propose Kan Extension Transformers (KETs) as a unifying categorical framework for a diverse group of Transformer implementations. The core claim is that a Transformer layer can be viewed as a weighted structured extension operator: standard attention is the singleton-neighborhood case, Geometric Transformer style incidence mixing is a sparse edge-restricted case, and KET is the higher-order simplicial case. This lens also clarifies a bridge to diffusion-style completion. When the extension operator acts on detached predictive carriers instead of teacher-forced hidden states, it becomes a valid self-conditioning mechanism that exposes noncausal structure without leaking gold future tokens. We include a comprehensive experimental validation of 12 different Transformer implementations varying across strict-causal and predict-detach regimes on Penn Treebank, WikiText-2, and WikiText-103. In the strict-causal setting, quadratic KET is the strongest model among the compared causal architectures on WikiText-2 and WikiText-103. Across all datasets, however, the largest gains come from the predict-detach regime rather than from changing the neighborhood family alone.

• The paper introduces a categorical formulation that unifies standard attention, diffusion, and higher-order (simplicial) aggregation via Kan extensions.
• It demonstrates that predict-detach self-conditioning enables noncausal neighborhood aggregation without violating causal constraints.
• Empirical results show quadratic KET strongly improves language modeling perplexity on larger datasets compared to baseline Transformers.

Kan Extension Transformers: A Categorical Synthesis of Attention, Diffusion, and Predict-Detach Self-Conditioning

Introduction and Categorical Framework

"Kan Extension Transformers: A Categorical Unification of Attention, Diffusion, and Predict-Detach Self-Conditioning" (2605.27259) advances the formal understanding and practical realization of neural LLMs by positing that Transformer layers can be viewed through the lens of Kan extensions—universal categorical constructions for extending functorial assignments. The core insight is that standard attention, geometric-inductive neighborhood mixing, and higher-order aggregation (involving simplices beyond edges—such as triangles or larger cliques) are all special cases of weighted extension operators. This categorical viewpoint enables a unified treatment of token-based attention (singleton neighborhoods), learned geometric adjacencies (Topocoend), and higher-order (simplicial) contexts, thereby linking attention and diffusion-style completion within a common formalism.

This framework does not simply repackage existing mechanisms; it establishes a precise bridge that specifies when architectural generalizations (such as noncausal neighborhoods for contextualization) are valid—particularly depending on whether values are sourced from teacher-forced future hidden states (which results in leakage) or from detached predictive carriers (which preserves causal validity by cutting the backward computational path).

Model Classes and Structural Generalization

The modeling spectrum articulated here varies predominantly by the structure of the neighborhood system over which aggregation occurs:

• Standard Attention: Aggregates over token positions, corresponding to singleton neighborhoods. This is the canonical Transformer update rule.
• TopoCoend: Aggregates over learned low-dimensional topological neighborhoods (constructed via latent representations and k-NN in the learned space).
• Kan Extension Transformer (KET): Aggregates over a simplicial complex, including 0-simplices (tokens), 1-simplices (edges), and higher-order structures (motifs, faces). Incidence-restricted KET considers only incident simplices, while quadratic KET aggregates globally over all simplices, with attention kernels defined on token-simplex pairs.

In strict-causal regimes, where only prefix information is accessible, quadratic KET consistently outperforms other causal architectures—most notably on WikiText-2 and WikiText-103—demonstrating the advantage of leveraging higher-order neighborhood information in a structured, principled way. However, on Penn Treebank (PTB), the baseline Transformer remains slightly stronger, attributable to the small scale and low redundancy of the dataset.

Predict-Detach Self-Conditioning and Modal Boundaries

A major theoretical and empirical contribution is the demarcation of three information regimes as they pertain to language modeling:

1. Strict-causal: Only prefix information flows (autoregressive decoding); future information is excluded in both structure and value.
2. Gold noncausal: Noncausal neighborhoods are admitted but values are sourced from hidden states computed with knowledge of the true future. This regime is invalid for standard language modeling due to information leakage, and as expected, test perplexity collapses to trivial values.
3. Predict-detach: Noncausal neighborhoods are permitted, but the values transported are detached predictive surrogates, generated from causal prefixes and with gradients blocked. This regime aligns with valid self-conditioning and supports noncausal structural aggregation without violating causal constraints.

The categorical semantics of the detach() operation are formalized as a modal boundary: forward information passes, but the backward (gradient) map is set to zero. This guarantees that self-conditioning via predictive surrogates (as opposed to direct teacher-forced signals) does not create illicit learning pathways and remains causally sound even in the presence of non-prefix aggregation.

Experimental Results: Causal and Self-Conditioned Performance

Comprehensive experiments evaluated twelve model configurations across three datasets. The strict-causal results demonstrated that:

• Quadratic KET is the strongest causal model on WikiText-2 and WikiText-103, outperforming both baseline Transformer and geometric transformer (GT) variants.
• Incidence-restricted KET, which reduces complexity by aggregating only over incident simplices, closely trails quadratic KET and provides an efficient trade-off.
• On smaller datasets, such as PTB, the baseline transformer matches or marginally exceeds the performance of KET, likely due to overfitting or limited higher-order structure.

In the predict-detach regime, the results are decisive:

• Across all datasets, transitioning from hidden-state carriers to detached predictive carriers yields an order-of-magnitude reduction in perplexity for KET variants—e.g., PTB: 133.37 → 31.43, WT2: 156.42 → 38.23, WT103: 210.30 → 51.89 for quadratic KET; even more pronounced for incidence-restricted versions.
• GT-PD (Geometric Transformer with predict-detach) achieves strikingly low perplexity, converging toward the regime where teacher-forced future information is present.
• TopoCoend only modestly improves under self-conditioning, indicating the principal gains stem from the information regime, not simply from increased neighborhood flexibility.

Figure 2: Validation perplexity trajectories across PTB, WikiText-2, and WikiText-103, showing that predict-detach self-conditioning (GT/KET-PD) produces steep gains across all architectures, with quadratic KET outperforming other strict-causal models on WikiText-2 and WikiText-103.

These results underline the dominance of information regime over architectural neighborhood complexity in determining empirical language modeling performance.

Connection to Blockwise Denoising and Structured Completion

Extending the categorical lens, the work explores structured completion tasks such as blockwise denoising—framing block prediction not as a one-shot mapping $C_t \to \Sigma^B$ but as a filling-in (completion) of a partially specified block $(C_t, \tilde{x}_{t+1:t+B}) \to x_{t+1:t+B}$. This perspective makes explicit the relation between left Kan extensions (aggregation/extension from evidence) and right Kan extensions (completion under constraints), with denoising falling into the latter.

Empirical findings indicate that denoising objectives are dramatically easier than direct block prediction, sharply reducing perplexity and depending primarily on the shift of information regime rather than backbone. The task is naturally interpreted simplicially: corrupted blocks correspond to horns (partial simplices), and denoising is horn-filling (completing coherent structures).

Theoretical and Practical Implications

The categorical reformulation of Transformer updates as Kan extensions clarifies both the scope and the limits of generalization when moving from sequence-local to higher-order, nonlocal, or learned geometric neighborhoods. The formalism provides tractable guidance for architectural design, indicating where further relaxation is empirically valid (predict-detach) and where leakage occurs (gold noncausal).

Practically, the results suggest that major empirical advances are available by rethinking the information regimes in language modeling—particularly by harnessing self-conditioned architectures that exploit richer noncausal neighborhoods while retaining causal learnability via gradient detachment. This insight is potentially transformative for future large-scale models, blockwise prediction, infilling, and other structured sequence modeling tasks.

On the theoretical side, the categorical synthesis unifies diverse literature on attention, diffusion, and geometric/topological neural models, and provides a principled backbone for extending results to multiplex graphs, hypergraphs, sheaf-based networks, and more abstract data domains.

Conclusion

Kan Extension Transformers (KETs) provide a categorical unification of Transformer-style attention, geometric/topological aggregation, and self-conditioning via predict-detach. Empirical results show that the chief performance gains in language modeling arise from switching the information regime rather than merely increasing the combinatorial or geometric complexity of the neighborhood system. The predict-detach paradigm enables the use of noncausal structure in a causally valid way, suggesting a broad spectrum of applications in structured completion, diffusion, and beyond. The categorical foundation established here offers a robust design space for future AI system architectures and invites exploration of richer, information-regime-based objectives for structured prediction.