🤖 machine learning

Journey Operators for Structured Multi-Axis Composition

This paper introduces a theoretical framework for structured multi-axis data composition using "journey operators" that generalize Rotary Position Embeddings to content-adaptive settings, establishes commutativity as the condition for path independence, and proposes the JoFormer model to leverage these inductive biases for improved performance in vision, language, and length generalization tasks.

Original authors: Mahesh Godavarti

Published 2026-07-30

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Mahesh Godavarti

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build a giant, magical library where every book knows exactly where it belongs on the shelf. In the world of artificial intelligence, specifically the kind that reads sentences or looks at pictures, the "books" are pieces of data like words or pixels. For a long time, computers have been really good at understanding what these pieces are (the content), but they've struggled a bit with where they are relative to each other. This is called "positional encoding." Think of it like a GPS for data: without it, the computer might know that "dog" and "man" are in a sentence, but it wouldn't know if the dog bit the man or the man bit the dog.

To solve this, scientists use math to give every piece of data a tiny, invisible "rotation" based on its spot. Imagine spinning a compass needle slightly for every step you take to the right, and a different spin for every step you take down. If you move right then down, the needle should end up in the same spot as if you moved down then right. This idea of moving through space and keeping track of direction is the foundation of the paper we are about to explore. It asks a simple but deep question: Can we make these "spins" work together perfectly so that the computer understands the structure of the world, whether it's a line of text, a grid of pixels, or a 3D volume?

The Journey of the Data

This paper, titled "Journey Operators for Structured Multi-Axis Composition," introduces a new way to think about how data moves and mixes inside an AI. The authors, led by Mahesh Godavarti, propose a framework where every piece of data carries not just its content (like the word "cat" or a red pixel), but also a tiny "transformation" or "rotation" for every direction it can move.

Let's use a playful analogy: Imagine every data point is a traveler carrying a backpack. In a standard AI model, the backpack just holds the traveler's stuff. But in this new framework, every traveler also carries a special "compass" for every direction they can walk (up, down, left, right). When two travelers meet to combine their stories (which is what an AI does when it processes data), the first traveler uses their compass to rotate the second traveler's backpack before they merge. This rotation changes how the second traveler's story is told, depending on where the first traveler is standing.

The paper calls this process a "journey." If you want to know how a piece of data at point A relates to a piece of data at point B, you calculate the "journey operator"—the total rotation you would get if you walked from A to B.

The Big Discovery: The "Flat" vs. "Curved" World

The most exciting part of the paper is the discovery of a rule that determines whether these journeys make sense. The authors found that for the AI to understand that "moving right then down" is the same as "moving down then right" (a property called path independence), the compasses (the rotations) must commute.

In math-speak, "commuting" means the order of operations doesn't matter. If you rotate a piece of paper 90 degrees to the right and then 90 degrees down, you get a different result than if you rotate it down first and then right—unless your rotations are special. The paper proves that if the rotations are "flat" (meaning they act on separate, non-interfering parts of the data, like rotating the X-axis and Y-axis independently), the journey is consistent. But if the rotations are "curved" (they interfere with each other), the journey depends on the exact path you took, which can confuse the AI.

The paper suggests that the best way to build these AI models is to stick to the "flat" regime. This explains why existing methods like RoPE (Rotary Position Embedding) work so well: they naturally use these commuting, flat rotations. The paper shows that if you try to use more complex, "curved" rotations (like dense, tangled rotations), you lose the ability to have a consistent sense of direction, and the positional structure collapses.

The New Model: JoFormer

Based on this theory, the authors designed a new model called JoFormer (Journey-based Transformer). The key innovation here is that JoFormer doesn't just use these rotations to help the AI score how much one word should pay attention to another (which is what older models did); it also uses the rotations to transform the actual data (the "values") before it is combined.

Think of it like this: In an old model, the AI might say, "I'm interested in the word 'dog' because it's near 'bark'." In JoFormer, the AI says, "I'm interested in the word 'dog', and I'm also going to spin the meaning of 'dog' based on how far away it is from me, so it fits perfectly into my current sentence."

The paper tests JoFormer on three different types of tasks:

Vision (Images): They tested it on CIFAR-100 and ImageNet. The results suggest that adding this "value rotation" (spinning the data) helps the model see better. For example, on CIFAR-100, a specific version of JoFormer improved accuracy by about 2% compared to standard methods. On ImageNet, it saw a smaller but consistent gain of 0.40%.
Language (Text): They tested it on a Wikipedia language model. The results suggest that the new model (specifically the "projected" version where the rotations depend on the content of the text) learned to predict the next word better, achieving lower "perplexity" (a measure of how surprised the model is) than the standard models.
Length Generalization: This is a cool test. They trained the model on short sentences (512 words) and asked it to read much longer ones (up to 4096 words). Standard models got very confused and their performance tanked. JoFormer, however, handled the long text much better, with its performance dropping only slightly (a ratio of 1.02x) compared to the standard model which dropped significantly (6.29x).

What the Paper Says It Doesn't Do

It's important to be clear about what this paper is not claiming. The authors are very careful to say that these results are "single-seed sanity checks." This means they ran the experiments once with a specific random starting point to see if the idea works in principle, not to claim they have built the absolute best AI in the world. They admit that their results are not "state-of-the-art" benchmark claims yet.

They also rule out certain ideas. For instance, they prove that if you try to use a full, complex rotation system (the full orthogonal group $O(d)$ ) instead of the simpler, flat rotations, the AI loses all sense of position. It's like trying to navigate a city where every street sign points in a different, confusing direction; the model just gives up on knowing where things are.

The Takeaway

In simple terms, this paper argues that the secret to making AI understand structure isn't just about adding more data or bigger computers; it's about giving the data a consistent, mathematical "compass" that works the same way no matter which way you turn. By ensuring these compasses don't fight each other (commuting), and by using them to actually spin the data as it moves through the network, the authors have created a model (JoFormer) that suggests a more robust way for AI to handle images, text, and long sequences.

While the paper doesn't claim to have solved everything, it provides a strong theoretical map showing why certain methods work and offers a new, promising path forward for building smarter, more structured AI systems. The experiments suggest that this approach has real, observable benefits, especially when dealing with long inputs or complex visual patterns, but the journey to fully understanding its potential is just beginning.

Technical Summary: Journey Operators for Structured Multi-Axis Composition

Problem Statement

Current machine learning architectures struggle to model data with structure along multiple axes (e.g., sequences, images, trees, volumetric data) in a way that aligns with vector-based learning mechanisms. While 1D compositional structures can be modeled using non-commutative semigroups or free groups, there is no broadly accepted algebraic framework for 2D or multi-dimensional composition. Existing algebraic systems often fail to support the specific needs of modern architectures like Transformers, particularly regarding how data composes and how relative positions are described across independent axes.

A critical tension exists: standard attention mechanisms aggregate values without non-commutative composition, effectively discarding the structural context of how data points relate along different axes. Furthermore, while Rotary Position Embedding (RoPE) successfully models 1D relative positions, extending this to multi-dimensional data requires a principled algebraic framework that ensures path independence (i.e., composing "right then down" yields the same result as "down then right") and supports content-adaptive inductive biases.

Methodology: The Algebraic Framework

The paper introduces a compositional framework where each data item carries both content ( $v \in \mathbb{R}^d$ ) and a set of axis-step generators ( $R_1, \dots, R_D \in G \le O(d)$ ), where $G$ is a group of norm-preserving orthogonal transformations.

Core Operation

In 1D, the composition of two data points $(a, A)$ and $(b, B)$ is defined as:
$(a, A) \circ (b, B) := (a + Ab, AB)$
Here, the content of the composite is the sum of the first point's content and the second point's content rotated into the first's frame. The axis-step generator of the composite is the product $AB$. This operation is associative but non-commutative.

In $D$ -dimensions, each datum carries one generator per axis. Composition along a specific axis $k$ applies that axis's generator to rotate the second point's content and accumulates the generator for that axis, leaving other axes unchanged.

The Journey Operator

A "journey" is defined as the path connecting two positions. The journey operator $P_{j \to k}$ is the product of axis-step generators along the path from position $j$ to $k$ .

Path Independence: The paper proves that composition and movement across axes are path-independent (i.e., the result depends only on displacement, not the order of axes) if and only if the per-axis axis-step generators commute ( $R_i R_j = R_j R_i$ ).
Flat vs. Curved Regimes: When generators commute, the framework operates in a "flat" regime where the journey operator is a pure function of displacement ( $P_{j \to k} = R^{k-j}$ ). When they do not commute, the regime is "curved," and relative transformations depend on absolute position.

Theoretical Classification

Under specific axioms (bilinearity, toral frame invariance, cocycle compositionality, and norm preservation), the paper classifies all compatible attention mechanisms. It proves that the only compatible mechanism is a block-diagonal $SO(2)^{d/2}$ rotation. This result explains why RoPE and its variants (2D-RoPE, tree PEs) arise naturally as special cases. Conversely, if the full orthogonal group $O(d)$ is assumed, positional structure collapses entirely.

JoFormer Architecture

Motivated by these findings, the authors propose JoFormer (Journey-based Transformer), which integrates the journey operator into the attention mechanism:

Score Side: Rotates queries and keys by their respective absolute operators ( $A_i q_i, A_j k_j$ ) before computing attention scores.
Value Side: Applies the journey operator to values before aggregation: $c_k = \sum_j \alpha_{kj} P_{j \to k} v_j$ . This brings compositional structure to the value side, which standard attention lacks.

The paper presents three variants:

JoFormer-fixed: Uses linear-in-position angles (standard RoPE frequencies) with value rotation.
JoFormer-learned: Uses learned frequency vectors per layer.
JoFormer-projected: Computes angles from the residual stream content via an MLP, creating a data-dependent journey operator.

Key Contributions

Compositional Framework: A unified algebraic framework where data points carry content and axis-step generators, defining composition as a single associative operation that extends naturally from 1D to $D$ -dimensions.
Theoretical Results:
- Proof that path-independent composition requires commuting axis-step generators.
- Classification showing that under natural axioms, the only compatible mechanism is block-diagonal $SO(2)^{d/2}$ rotation, recovering RoPE as a special case.
- Proof that in the flat regime, value transport via the journey operator is translation equivariant.
- Demonstration that the Discrete Fourier Transform (DFT) is a special case of this compositional embedding.
JoFormer Architecture: A new model design that applies the journey operator to values, unifying standard attention, State-Space Models (SSMs), and JoFormer under a "value-path formula."
Value-Path Hierarchy: A theoretical bridge showing that SSM recurrences unroll into the same value-path formula as attention, clarifying the structural differences between these families.

Experimental Results

The authors conducted single-seed sanity checks across vision, language, and length generalization tasks to test the framework's hypotheses:

Vision (CIFAR-100 & ImageNet):
- On CIFAR-100, applying Value (V) rotation (rotating values and inverse-rotating the output) consistently improved accuracy across different model sizes and frequency types. The synergy between V rotation and learnable frequencies yielded a +1.88% gain over standard 2D-RoPE.
- On ImageNet-1K (ViT-S), JoFormer with V rotation achieved a +0.40% Top-1 accuracy improvement over RoPE2D.
- Experiments confirmed the "flat regime" hypothesis: dense (curved) rotations (LieRE) performed worse than axial (commuting) rotations.
Language Modeling (Wikipedia):
- JoFormer-projected (content-dependent angles) achieved the lowest perplexity across all tested configurations compared to RoFormer and fixed/learned JoFormer variants.
Length Generalization:
- JoFormer-projected demonstrated superior extrapolation capabilities. When trained on sequence length 512 and tested on 4096, it achieved a perplexity ratio of 1.02x, significantly outperforming RoPE (6.29x) and JoFormer-fixed (4.41x). This suggests that data-dependent angles are crucial for length generalization beyond just V rotation.

Significance and Claims

The paper claims to provide a foundational algebraic framework that resolves the tension between multi-axis compositional structure and vector-based learning. Its significance lies in:

Theoretical Unification: It explains why RoPE-like methods are natural outcomes of specific symmetry constraints and provides a common notation for attention, SSMs, and JoFormer.
Inductive Bias: It identifies a missing structural bias in standard attention: the lack of non-commutative composition on the value side. JoFormer fills this gap.
Practical Utility: Initial experiments suggest that these inductive biases (specifically V rotation and content-dependent angles) yield observable improvements in vision tasks and length generalization.

The authors remain modest, noting that their experiments are single-seed sanity checks rather than definitive benchmark claims, and that the framework's value is primarily in providing a principled design rule for positional mechanisms. They explicitly state that while the theory motivates the architecture, it does not prove that the specific variants (like JoFormer-projected) will always outperform baselines, but rather that they are consistent with the derived value-path hierarchy.