Manifold-Based Token Interaction Redesign: From Fixed Windows to Learned Flow Fields

TL;DR: Instead of using fixed local windows for Grassmann mixing, can we learn adaptive flow fields over token sequences, inspired by manifold-guided GNNs and cross-token mixing in transformers? Try a model where the “window” for subspace formation is dynamically selected based on learned manifold-based affinities or structural cues.

Research Question: Can adaptive, data-driven selection of token interactions—using learned geometric flow fields—outperform fixed-window Grassmann flows or standard attention in modeling complex dependencies?

Hypothesis: Allowing the model to learn which tokens should interact geometrically (rather than forcing locality) will result in better long-range reasoning and more flexible context modeling, closing the gap to transformers in tasks requiring global dependencies.

Experiment Plan: Implement a geometric token interaction mechanism where subspaces are formed over dynamically selected token groups, guided by learned affinities (possibly via a lightweight GNN or manifold kernel). Evaluate on tasks with varying dependency lengths (e.g., synthetic copy tasks, long-context language modeling). Analyze the learned interaction graphs to identify emergent structures and compare efficiency to full attention.

References:

• Li, X., Li, C., Wei, X., & Yang, F. (2021). Manifold Guided Graph Neural Networks for Skeleton-based Action Recognition in Human Computer Interaction Videos. International Conference on Signal Processing and Machine Learning.
• Wang, L., Zheng, Z., Kumar, N., Wang, C., Guo, F., & Zhang, P. (2024). Multilevel Class Token Transformer With Cross TokenMixer for Hyperspectral Images Classification. IEEE Transactions on Geoscience and Remote Sensing.
• [Original paper]