Towered Graph: Kan Extension Transformers

• Towered Graph is a framework that redefines Transformer layers as Kan extensions, interpreting attention as weighted aggregation over tokens and higher-order structures.
• It unifies standard self-attention with geometric mixing (TopoCoend) and simplicial aggregation, thereby capturing complex relational patterns in structured data.
• The approach leverages a predict-detach self-conditioning mechanism and a diffusion-style completion strategy to significantly improve performance in language modeling benchmarks.

Kan Extension Transformers (KETs) reframe a Transformer layer as a weighted structured extension operator grounded in Kan extensions from category theory. The core idea is to choose a source neighborhood system of objects (tokens, edges, higher-order simplices, or learned geometric neighbors), attach values (“carriers”) to those objects, and transport them to target token positions via weights determined by attention scores, learned topology, or simplicial incidence. This view unifies standard attention (singleton token neighborhoods), Geometric Transformer incidence mixing (sparse edge-restricted neighborhoods), and KET’s higher-order simplicial aggregation. It also clarifies a bridge to diffusion-style completion and a valid self-conditioning mechanism based on detached predictive carriers.

Intuition and categorical background: Kan extensions and weighted structured extension

• The paper organizes source values and target structure functorially and uses Kan extensions as universal extension operators: left Kan for aggregation from local evidence and right Kan for completion under constraints. The formal definitions and pointwise formulas used in the paper are:

Definition (left and right Kan extensions): Let

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

be functors. We would like to construct a functor

$H : \mathcal{D} \to \mathcal{E}$

that extends $F$ along $K$.

\begin{definition} The left Kan extension of $F$ along $K$, denoted $\operatorname{Lan}_K F$, is a functor

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

together with a natural transformation

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

such that for any other pair

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

there exists a unique natural transformation

$H : \mathcal{D} \to \mathcal{E}$0

with

$H : \mathcal{D} \to \mathcal{E}$1

Dually, the right Kan extension $H : \mathcal{D} \to \mathcal{E}$2 is the universal extension characterized by maps

$H : \mathcal{D} \to \mathcal{E}$3

\end{definition}

Pointwise formulas via colimits and limits: \begin{equation} (\operatorname{Lan}K F)(d) \;\cong\; \operatorname*{colim}{(K \downarrow d)} F. \label{eq:pointwise_lan} \end{equation} \begin{equation} (\operatorname{Ran}K F)(d) \;\cong\; \operatorname*{lim}{(d \downarrow K)} F. \label{eq:pointwise_ran} \end{equation}

Enriched coend formulations (weighted aggregation for left Kan; compatibility for right Kan): \begin{equation} (\operatorname{Lan}K F)(d) \;\cong\; \int^{c \in \mathcal{C} \mathcal{D}(Kc,d)\,\otimes\,F(c). \label{eq:kan_coend} \end{equation} \begin{equation} (\operatorname{Ran}_K F)(d) \;\cong\; \int{c \in \mathcal{C} [\mathcal{D}(d,Kc),\,F(c)], \label{eq:ran_end} \end{equation}

Common setup across model families: Let

$H : \mathcal{D} \to \mathcal{E}$4

be target token positions and $H : \mathcal{D} \to \mathcal{E}$5 a source neighborhood system (tokens, edges, simplices, or learned geometric neighbors). The paper introduces a carrier assignment

$H : \mathcal{D} \to \mathcal{E}$6

and weights $H : \mathcal{D} \to \mathcal{E}$7 that describe how source object $H : \mathcal{D} \to \mathcal{E}$8 contributes to target token $H : \mathcal{D} \to \mathcal{E}$9. The shared computational form is: \begin{equation} h't \;\approx\; \int^{\sigma \in \mathcal{N} W(t,\sigma)\otimes X(\sigma) \;\approx\; \sum{\sigma \in \mathcal{N} w(t,\sigma)\,V(\sigma), \label{eq:generic_kan_pool} \end{equation} where the coend notation emphasizes structured weighted extension rather than an arbitrary sum. In practice, this becomes a familiar “weighted pooling” update. Terminology: “carriers” are the values attached to source objects (e.g., token hidden states, edge embeddings, or detached predictive embeddings), “neighborhoods” are the families of source objects over which we aggregate, and weights are constructed either from attention kernels (softmax of Q·K), topological proximity in a learned latent space, or simplicial incidence restrictions.

Attention as the singleton-neighborhood case

• Standard self-attention is recovered by taking the source neighborhood system to be tokens only: $F$0 Each source object is a token $F$1, and the update at target token $F$2 is: $F$3 Categorically, this is a weighted left-Kan-style extension over token objects, with attention scores providing the hom-weights. The paper maps attention’s Q, K, V into the weighting: the KET quadratic form below explicitly uses

$F$4

so when $F$5 ranges only over tokens, this reduces to standard token-to-token attention.

Geometric and higher-order cases: TopoCoend and KET

• Topological neighborhoods (TopoCoend): The source objects remain tokens, but neighborhoods are learned in a latent space via

$F$6

and a fuzzy k-NN graph is constructed over the $F$7’s. The TopoCoend update is \begin{equation}

h'_t

h_t + \sum_s w_{\mathrm{topo}(t,s)\,V_s. \label{eq:topocoend_update} \end{equation} Weights $F$8 are induced by topological proximity rather than sequence position alone.

• Simplicial neighborhoods (KET): KET changes the source objects themselves:
  • 0-simplices: tokens $F$9,
  • 1-simplices: edges such as $K$0,
  • optional higher simplices: faces, motifs, or larger spans.

Let $K$1 be the simplicial indexing category. Each simplex $K$2 has value $K$3, weighted by $K$4. The KET update: \begin{equation}

h'_t

h_t + \sum_{\sigma \in \mathcal{N}{\mathrm{simp} w{\mathrm{simp}(t,\sigma)\,V(\sigma). \label{eq:simplicial_update} \end{equation}

Two realizations:

1. Quadratic KET (global pooling over all simplices): \begin{equation}

   h'_t

   h_t + \sum_{\sigma \in \mathcal{N}{\mathrm{simp} w(t,\sigma)\,V(\sigma), \qquad w(t,\sigma)=\operatorname{softmax}(Q_t^\top K\sigma). \label{eq:quadratic_ket} \end{equation} This behaves like attention from tokens to simplices, typically $K$5.

2. Incidence-restricted KET (sparse local mixing over incident simplices): \begin{equation} h't \;\approx\; h_t + \sum{\sigma \ni t} \phi\bigl(V(\sigma)\bigr). \label{eq:incidence_ket} \end{equation} Edge-only example: \begin{align} e_t &= \psi([v_{t-1},v_t]), \ h'_t &= h_t + \phi(e_t) \qquad \text{(causal incidence)}. \end{align} This reduces complexity to $K$6 and makes the Geometric Transformer (GT) connection transparent.

Relationship to Geometric Transformers:

• The geometric branch of a GT is an incidence-restricted Kan-style update with local convolution/message-passing maps, placing models in a hierarchy:
  • Attention: weighted extension over tokens.
  • TopoCoend: weighted extension over learned topological neighborhoods.
  • Incidence-restricted KET: weighted extension over incident simplices.
  • Quadratic KET: weighted extension over all simplices with a learned global kernel.
  • GT: efficient incidence-restricted special case of KET.

Note: The paper does not introduce formulas for boundary/coboundary operators ($K$7), Hodge Laplacians, or weighted cochain complexes; higher-order structure is captured through simplicial neighborhoods and incidence, not explicit homological operators.

Bridge to diffusion-style completion: right-Kan intuition and block denoising

• The paper connects KET’s left-Kan aggregation perspective to diffusion-style completion via a right-Kan intuition: completion seeks values compatible with multiple local constraints. It implements block denoising (structured completion), contrasting it with direct block prediction:

Direct block prediction: \begin{equation} f : C_t \longrightarrow \Sigma^B, \label{eq:block_product_map} \end{equation} with

$K$8

Denoising block completion: \begin{equation} f : (C_t,\tilde{x}{t+1:t+B}) \longrightarrow x{t+1:t+B}, %\label{eq:denoise_map} \end{equation} where $K$9 is a corrupted block. Corruption schedule: $F$0 ranging from 0.05 to 0.50. The completion task is framed as “horn filling” in a simplicial analogy (partial simplex to full simplex). The paper does not provide explicit Laplacian update equations (e.g., $F$1); the connection is conceptual via right Kan (limits/ends) as compatibility-based completion.

Predict-detach self-conditioning: detached predictive carriers and causal validity

• KET’s most important modal mechanism is predict-detach: transporting detached predictive carriers through (possibly noncausal) neighborhoods to expose structure without gold-future leakage.

Predictive carrier construction: \begin{equation}

\hat e_t

detach!\left(softmax(\ell_t / T)\, E\right), \label{eq:predict_detach} \end{equation} where $F$2 is the embedding matrix, $F$3 is a temperature, and $F$4 are next-token logits from the causal hidden state. The same formula appears in the appendix: \begin{equation}

\hat e_t

\operatorname{detach}!\Bigl( \operatorname{softmax}(\ell_t/T)\,E \Bigr), \qquad \ell_t = W_o h_t. \label{eq:appendix_pred_detach} \end{equation} Why detach matters:

• No gold-future values: carriers derive from prefix-valid predictive states, not teacher-forced future hidden states.
• No leakage gradient: detach blocks the auxiliary branch from backpropagating targets into carriers later consumed by noncausal aggregation.

Information regimes:

• Strict-causal: neighborhoods and values restricted to prefix-valid information.
• Gold noncausal: invalid leakage regime; mixing teacher-forced future hidden states collapses perplexity toward 1.
• Predict-detach: noncausal neighborhoods allowed, but transported values are detached predictions from causal states. This modality preserves causal validity.

Computational regimes and trade-offs

• Quadratic KET: global simplex pooling with attention-like weights, complexity typically $F$5; benefits most from GPU acceleration.
• Incidence-restricted KET: only incident simplices contribute, typically $F$6; efficient sparse approximation to quadratic neighborhoods.
• TopoCoend: learned geometric neighborhoods; sparse and adaptive depending on k-NN graph size.

Experimental setup and results

• Datasets and base settings: Penn Treebank (PTB), WikiText-2 (WT2), WikiText-103 (WT103). Context length $F$7, $F$8 layers, $F$9 hidden dimension (unless otherwise stated). Evaluation metric: perplexity.

Strict-causal comparison (Table 1 in paper): \begin{tabular}{lccc} 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 \ \end{tabular} Key takeaways:

• On PTB, the plain Transformer remains strongest among strict-causal models.
• On WT2 and WT103, quadratic KET (KET-Quad-C) is strongest among compared causal architectures.
• Incidence-restricted KET is slightly weaker but close on larger datasets, consistent with efficient sparse approximation.

Predict-detach ablation (Table 2 in paper): \begin{tabular}{lccc} Model & PTB & WT2 & WT103 \ KET-Quad-C & 133.37 & 156.42 & 210.30 \ KET-Quad-PD & 31.43 & 38.23 & 51.89 \ KET-Inc-C & 137.19 & 161.12 & 213.76 \ KET-Inc-PD & 6.54 & 19.08 & 47.17 \ GT-PD & 1.05 & 1.59 & 12.84 \ \end{tabular} Interpretation:

• Replacing hidden-state carriers with detached predictive carriers dramatically improves all KET variants: quadratic KET drops from 133.37→31.43 (PTB), 156.42→38.23 (WT2), 210.30→51.89 (WT103).
• Incidence-restricted KET improves even more on PTB and WT2.
• GT-PD is the strongest self-conditioned baseline in this comparison.
• The dominant empirical effect comes from the information regime itself (predict-detach), not just changing the neighborhood family.

Block denoising experiments (structured completion; Table 3 in paper): context $K$0, $K$1, block size $K$2, depth sweeps $K$3 and $K$4.

• Denoising is dramatically easier than direct block prediction across datasets and depths; direct block prediction improves with depth, but denoising remains low perplexity with modest backbone effects.
• This echoes the right-Kan completion intuition: providing partial future structure (corrupted block) makes completion easier than one-shot generation of $K$5.

Training and implementation details:

• Main causal language-model comparisons:
  • Optimizer: AdamW
  • Learning rate: $K$6
  • Weight decay: $K$7
  • Batch size: $K$8
  • Context length: $K$9
  • Layers $\operatorname{Lan}_K F$0, hidden dimension $\operatorname{Lan}_K F$1, attention heads $\operatorname{Lan}_K F$2
  • Training iterations: $\operatorname{Lan}_K F$3
  • Evaluation cadence: every $\operatorname{Lan}_K F$4 steps
  • Self-conditioned runs: predictive carriers with temperature $\operatorname{Lan}_K F$5
  • TopoCoend: $\operatorname{Lan}_K F$6, $\operatorname{Lan}_K F$7
• Structured completion runs:
  • Matched Transformer and incidence-KET backbones
  • Optimizer: AdamW
  • Learning rate: $\operatorname{Lan}_K F$8
  • Weight decay: $\operatorname{Lan}_K F$9
  • Context length: $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$0
  • Hidden dimension: $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$1
  • Layers: $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$2 and $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$3
  • Block size: $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$4
  • Training steps: $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$5
  • Evaluation cadence: every $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$6 steps
  • Denoising steps: $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$7
  • Gradients clipped to norm $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$8
  • Seeds recorded; hardware includes Apple MPS and NVIDIA DGX Spark; quadratic KET benefits disproportionately from GPU acceleration.

Layer-level algorithms (from the paper’s pseudocode):

• Quadratic Kan Extension Block: 1) Construct simplex values $$\operatorname{Lan}_K F : \mathcal{D} \to \mathcal{E}$$9 for all $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$0 in the simplex set $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$1. 2) Construct keys $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$2. 3) For each token $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$3: compute $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$4; select $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$5 by causal mask if enabled; compute weights $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$6; aggregate $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$7; update $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$8.
• Incidence-Restricted Kan Block (Edge-only): 1) For $$\eta : F \Rightarrow (\operatorname{Lan}_K F)\circ K$$9 to $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$0, form edge embeddings: $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$1. 2) For each token $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$2: accumulate $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$3 from incident edges (causal: $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$4; noncausal variants may add $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$5); update $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$6.
• Predict-detach carriers: $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$7 with $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$8. Use $$G : \mathcal{D} \to \mathcal{E}, \qquad \gamma : F \Rightarrow G\circ K,$$9 as the value base $H : \mathcal{D} \to \mathcal{E}$00 for constructing edges/simplices or for TopoCoend embedding. Detach removes gradient flow through the predictive branch while preserving forward carriers.

Leakage diagnostic:

• Shuffle target tokens within a batch and check whether perplexity collapses. Gold noncausal regimes collapse toward 1 under shuffling (invalid). Predict-detach runs stay sensible, indicating no gold-future leakage.

Related work and distinctions

• The paper’s contribution is unifying attention, geometric mixing, learned topological neighborhoods, and self-conditioned completion within the categorical language of Kan extensions:
  • Attention is a weighted left-Kan-style extension over tokens.
  • Geometric Transformers are incidence-restricted Kan updates over local simplicial structure.
  • TopoCoend is a coend-style weighted aggregation over learned geometric neighborhoods.
  • Predict-detach provides a modal boundary enabling noncausal structural forward computation without invalid gradient paths, aligning with denoising/diffusion intuitions (right-Kan completion).
• The paper situates KET relative to: topological/simplicial attention (e.g., GSAN, Simplicial Attention Networks), sheaf neural networks and diffusion, hypergraph Transformers, and diffusion models (conceptual link rather than score-based training). What is novel: treating layers explicitly as Kan extensions; lifting neighborhood families to simplicial objects; and formalizing predict-detach as a valid self-conditioning regime.

Practical guidance, limitations, and future directions

• When to use KETs:
  • If higher-order structure beyond token adjacency is important (e.g., motifs, spans, edges, faces), KET’s simplicial neighborhoods can improve expressivity.
  • When computational budget allows, quadratic KET provides strong causal baselines on larger corpora (WT2, WT103).
  • For large gains, focus on information regime: predict-detach self-conditioning yields the biggest improvements across datasets.
• Expected benefits beyond language modeling:
  • Any task that can be framed as structured extension or completion over a neighborhood system (e.g., structured prediction, infilling, blockwise denoising, nontext modalities with topological relations).
• Limitations:
  • Study scale is modest and centered on perplexity; results are single-seed without error bars.
  • The largest empirical gains arise from regime changes (predict-detach) rather than purely from changing neighborhood families; architectural conclusions should be made within regime.
  • Richer homological operators (boundary/coboundary, Hodge Laplacians) are not used in the reported KET implementations.
• Open questions/future directions:
  • Scale neighborhood design and information regimes separately.
  • Explore richer simplicial families, adaptive simplex selection, longer contexts with sparse approximations.
  • Hybrid systems: combine KET’s combinatorial structure with TopoCoend’s learned geometry.
  • Extend structured completion beyond small block denoising: infilling, partial-information objectives, and other right-Kan-style tasks.

Summary KETs formalize Transformer layers as weighted structured extensions via Kan extensions. Attention is the singleton case; Geometric Transformers are incidence-restricted sparse extensions; KET aggregates over higher-order simplices. This categorical viewpoint clarifies a bridge to diffusion-style completion (right-Kan intuition) and motivates predict-detach self-conditioning as a valid mechanism to exploit noncausal neighborhoods without leakage. Empirically, quadratic KET is strongest among causal models on WT2 and WT103, but the dominant improvements arise from the predict-detach regime itself. The paper provides algorithms for quadratic and incidence-restricted Kan blocks, training details, leakage diagnostics, and structured completion results, framing a broader design space where neighborhood systems and information regimes can be varied systematically.