Nestohedral Adaptive Token Tree

(1) Integration point: insert this module before two or three Transformer blocks and operate on the current token embeddings h_1,...,h_n. Build a sparse hypergraph H from token adjacency: connect neighboring positions for text, neighboring patches for images, and optionally retrieval-neighbor tokens. Maintain a rooted construct with singleton leaves and connected internal groups. Each active node is represented by its mean-pooled embedding plus a learned size embedding. Attention is performed among active nodes, and each node output is broadcast to its descendant leaves.

(2) Pseudocode:

H = build_local_hypergraph(n, radius=r)
T = singleton_tree(H)
for layer in selected_layers:
    hX = mean(h[i] for i in leaves(X)) for every active node X
    q, k, v = linear(hX)
    z = softmax(q @ k.T / sqrt(d) + tree_mask(T)) @ v
    split_score[X] = MLP([hX, size(X)])
    fusion_score[Y,X] = MLP([hY, hX, hY-hX]) for adjacent nodes
    contract highest-scoring valid edges if node count exceeds budget
    split valid nodes whose split score exceeds threshold
    broadcast z to descendant leaves

A contraction implements Y--X to X union Y; splitting is its inverse when both child labels are valid tubes and the recursive component condition holds. Use straight-through binary gates or Gumbel-Softmax during training, then deterministic thresholds during inference.

(3) Compute from the paper: tube validity, connected components, recursive tree construction, edge contractions, and the dimension statistic d(T). Estimate empirically: split and fusion scores, the target node budget, and possibly the hypergraph if it is learned. For a cheap MVP, use interval groups on sequences, where connectivity is trivial and every group is a contiguous span.

(4) First experiment: train a 6-layer, 128-dimensional Transformer on WikiText-2 or a synthetic long-range copy task. Compare dense attention and fixed two-stage token pooling at equal baseline quality. Sweep active-node budgets of 25%, 50%, and 75%. Measure validation perplexity, attention FLOPs, peak activation memory, and where splits occur. Success means at least 1.5x lower attention FLOPs at matched perplexity, with splits concentrated around punctuation, entity boundaries, or long-range dependency locations rather than uniformly.