Spacetime Optimal-Transport Attention (SO-TA)

• The paper introduces SO-TA, a tri-modal fusion module that replaces softmax attention with entropy-regularized optimal transport for aligning force-pose-derived queries with visual patches.
• It integrates synchronized RGB, force/torque, and proprioceptive pose streams with a diffusion-based sequence policy to handle contact-rich manipulation tasks under variable conditions.
• Empirical results demonstrate improved robustness and higher success rates in tasks like tight peg-in-hole assembly, outperforming traditional cross-attention and concatenation methods.

Spacetime Optimal-Transport Attention (SO-TA) is a tri-modal fusion module for contact-rich imitation learning that replaces softmax-normalized patch attention by an entropy-regularized Optimal Transport (OT) alignment between force-pose-derived sub-queries and visual patches. In the formulation introduced for visuo-haptic manipulation, SO-TA is paired with a diffusion-based sequence policy that maps synchronized windows of RGB, force/torque, and proprioceptive pose to future pose-action chunks. Its central claim is that explicit marginal constraints provide a structured inductive bias for conditioning-aware spatial selection, which is reported to be stable across illumination, distractors, and partial occlusion (Feng et al., 19 May 2026).

1. Definition and motivating problem

SO-TA was introduced for contact-rich manipulation tasks such as tight-clearance insertion, connector mating, polishing, and surface-conforming wiping, where discontinuous contact dynamics, partial observability, and strict safety constraints make fusion design consequential (Feng et al., 19 May 2026). The motivating observation is phase dependence: vision supplies global context before contact, force/torque feedback governs interaction after contact, and proprioceptive pose provides a consistent kinematic backbone. In this setting, standard cross-attention can learn correlations, but its softmax normalization imposes no explicit constraint on how much attention mass should go to different visual patches, while concatenation discards spatial structure entirely.

The SO-TA formulation therefore replaces the patch-selection step with entropy-regularized OT. Spatial aggregation is no longer only a learned similarity score; it becomes a constrained alignment between force-pose-conditioned “supply” and visual patch “capacity.” The official formulation is specific: SO-TA is not the policy itself, but the structured multimodal conditioning front-end for a diffusion policy (Feng et al., 19 May 2026).

A recurring misconception is to read “spacetime” as implying that the OT problem itself is solved jointly across temporal trajectories. In the reported architecture, the OT alignment is solved for each frame sample after merging batch and time dimensions, while time enters through observation windows of length $T_w$, a per-frame Transformer, a temporal Transformer, and a diffusion-based sequence model (Feng et al., 19 May 2026). This suggests that the “spacetime” designation reflects temporally contextualized multimodal fusion rather than a multi-frame OT coupling.

2. Task formulation, sensing streams, and sequence structure

The learning problem is defined over synchronized windows of RGB, force/torque, and proprioceptive pose. At each time index $t$, the observation tuple is a fixed-length window of length $T_w$ mapped to a future pose-action chunk of horizon $T_h$, and the policy follows Diffusion Policy with the action chunk as the diffusion variable and the fused observation encoding as the conditioning variable (Feng et al., 19 May 2026).

The architecture operates on a window of length $T_w$ and produces a fused encoding

$$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$

The image stream is center-cropped, resized to $224\times 224$, and encoded per frame by a ResNet-18 with GroupNorm and deterministic 2D positional encoding. Force/torque observations are represented as a 9-dimensional force channel after magnitude/direction splitting and normalization. Proprioceptive pose is represented as a 10-dimensional pose vector using a 6D rotation representation and incremental motion normalization (Feng et al., 19 May 2026).

The force and pose streams are mapped by learnable MLP encoders into token sequences

$$\mathbf{Z}^{\mathrm{force}},\mathbf{Z}^{\mathrm{pose}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}},$$

then concatenated and projected into a conditioning embedding

$$\mathbf{z}^{\mathrm{fp}}\in\mathbb{R}^{B\times T_w\times d_{\mathrm{att}}}.$$

Vision is encoded separately into patch tokens

$$\mathbf{Z}^{\mathrm{img}}\in\mathbb{R}^{B\times T_w\times N_{\mathrm{patch}}\times d_{\text{model}}},$$

with $t$0 from a $t$1 patch grid (Feng et al., 19 May 2026).

The experimental system uses $t$2, $t$3, $t$4, $t$5, $t$6, $t$7, OT temperature $t$8, and $t$9 Sinkhorn iterations. Diffusion training uses $T_w$0 steps and $T_w$1 denoising steps in the main setup. Online inference runs at $T_w$2 Hz (Feng et al., 19 May 2026).

3. Optimal-transport attention mechanism

SO-TA uses the force-pose embedding to form sub-queries and to define OT supply. After merging batch and time dimensions, for each frame sample $T_w$3,

$T_w$4

where $T_w$5. A second head predicts a row-marginal supply

$T_w$6

This means that the haptic/proprioceptive state determines how much transport mass each sub-query owns at that time step (Feng et al., 19 May 2026).

Each visual patch is projected to the OT space through key and value projections. The OT cost between sub-query $T_w$7 and patch $T_w$8 is the negative cosine similarity,

$T_w$9

With uniform patch capacity

$T_h$0

the transport plan solves

$T_h$1

subject to

$T_h$2

where

$T_h$3

The Sinkhorn algorithm is used in log-domain, with dual updates

$T_h$4

$T_h$5

and the plan is recovered as

$T_h$6

After solving OT, each sub-query aggregates visual values,

$T_h$7

SO-TA then performs a tied merge using the row marginals: $T_h$8 with

$T_h$9

The paper explicitly states why this tie matters: because the patch capacity $T_w$0 is uniform, a merge independent of $T_w$1 would collapse toward scaled mean pooling and lose conditioning dependence; tying the merge weights to $T_w$2 preserves force-pose conditioning without extra parameters (Feng et al., 19 May 2026).

This is the core formal difference from the matched-capacity cross-attention baseline, which replaces OT with standard scaled dot-product attention. In SO-TA, marginal constraints make the alignment jointly normalized across sub-queries and patches rather than only row-wise softmax-normalized (Feng et al., 19 May 2026).

4. Integration with diffusion policy and sequence modeling

SO-TA is inserted into a larger sequence-modeling stack. At each time step, a length-3 token sequence $T_w$3 is processed by a lightweight per-frame Transformer shared across frames. A pooled head produces a softmax gate $T_w$4, which combines the contextualized tokens into a frame-wise feature $T_w$5. Sinusoidal time encoding is then added, and a temporal Transformer produces the final fused representation $T_w$6 (Feng et al., 19 May 2026).

The diffusion policy uses the action chunk as the diffusion variable $T_w$7 and the fused observation encoding as the conditioning variable $T_w$8. Training follows the standard $T_w$9-prediction objective,

$$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$0

Online, a short reverse diffusion process produces a pose-action chunk, which is then used in receding-horizon control (Feng et al., 19 May 2026).

Training is done with AdamW, cosine warm-up, mixed precision, and EMA weights, on a single NVIDIA RTX 6000 Pro for 4,000 epochs. The matched-capacity SO-TA and cross-attention models each have 61M total parameters, decomposed into 40M diffusion and 21M fusion, while concatenation has 81M total (Feng et al., 19 May 2026).

A second misconception is that SO-TA is designed as a generic replacement for all temporal attention in sequence models. The reported implementation is narrower: OT is applied at the force-pose-to-patch fusion interface, then the resulting fused token participates in conventional sequence modeling through Transformers and diffusion. The paper therefore treats OT as a structured front-end inductive bias rather than a universal end-to-end OT backbone (Feng et al., 19 May 2026).

5. Empirical evaluation, robustness, and diagnostics

The reported evaluation covers three real-robot tasks using a bilateral teleoperation setup in which the operator sees the scene only through the wrist-mounted RGB stream and leader-side haptic feedback, matching deployment sensing; only follower-side signals are logged (Feng et al., 19 May 2026).

Tight peg-in-hole assembly uses a clearance below $$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$1 mm and is highly contact-sensitive. The dataset contains 198 teleoperated episodes, 32,974 synchronized tri-modal steps, about 3,297.4 s total, yielding 29,608 window-chunk pairs. Over about 200 rollouts, SO-TA reaches $$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$2 success, concatenation reaches $$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$3 with one failure, and cross-attention reaches $$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$4. Cross-attention also shows the heaviest completion-time tail (Feng et al., 19 May 2026).

BCM wiring-connector insertion is more visually cluttered and lower contrast than peg-in-hole. The dataset contains 396 demonstrations, 29,126 steps, about 2,916.2 s, yielding 22,430 training pairs. Success requires the connector to be fully seated under $$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$5 N downward force. The SO-TA policy achieves a mean success-only time close to human teleoperation, around 10 s, but with a long tail of 20–100 s and 7 timeouts. A cutoff-and-reset wrapper with a cutoff of 12 s is selected to minimize the 99th percentile total time-to-success while maintaining at least 99.5% success; with this wrapper, the 99th percentile drops to 47.7 s versus 70.85 s without cutoff, and simulated success reaches 100% over 1,000 episodes (Feng et al., 19 May 2026).

Curved-surface mark erasing requires localization on a curved surface, contact regulation, and sustained sliding friction. The dataset contains 199 demonstrations, 35,374 steps, about 3,437.4 s, yielding 31,991 training pairs. SO-TA attains a mean success-only time of 19.3 s versus 17.6 s for human teleoperation, with 166/200 successes under the strict all-marks-removed criterion. Failures correlate with under-represented mark locations near the boundary of the allowed region (Feng et al., 19 May 2026).

Robustness under visual perturbations is one of the clearest reported results. Under illumination changes, distractor pegs, and partial occlusion, SO-TA retains $$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$6 success $$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$7, while concatenation drops to $$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$8 $$\mathbf{Z}^{\mathrm{fused}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}}.$$9 (Feng et al., 19 May 2026). The modality ablations on the concatenation backbone further report that, over roughly 200 rollouts per variant under a 5 N safety threshold and 100 s timeout, the full tri-modal input gives about $224\times 224$0 success, removing pose gives $224\times 224$1, removing force/torque gives $224\times 224$2, and vision-only gives $224\times 224$3. Failures with zero-masked force/torque are dominated by exceeding the safety threshold, while failures with force/torque present are more often covariate-shift drift (Feng et al., 19 May 2026).

The diffusion-inference latency study reports average inference times from 9.6 ms to 16.2 ms for 1–10 denoising steps, corresponding to effective rates from 105 Hz to 62 Hz. After the second iteration, the first-step $224\times 224$4-error collapses sharply, and the first-step $224\times 224$5 MAE is already around $224\times 224$6 m for most settings beyond 2–3 steps. This is used to support the claim that the receding-horizon policy can meet the 10 Hz control deadline (Feng et al., 19 May 2026).

Interpretability is analyzed with OT-derived patch heatmaps and leave-one-out modality influence ratios. The leave-one-out ratio is

$224\times 224$7

In peg-in-hole, before contact, vision has the highest share at $224\times 224$8, with pose at $224\times 224$9 and force at $$\mathbf{Z}^{\mathrm{force}},\mathbf{Z}^{\mathrm{pose}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}},$$0. During hole searching and insertion, visual share drops to $$\mathbf{Z}^{\mathrm{force}},\mathbf{Z}^{\mathrm{pose}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}},$$1 and $$\mathbf{Z}^{\mathrm{force}},\mathbf{Z}^{\mathrm{pose}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}},$$2, while pose rises to $$\mathbf{Z}^{\mathrm{force}},\mathbf{Z}^{\mathrm{pose}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}},$$3 and force to $$\mathbf{Z}^{\mathrm{force}},\mathbf{Z}^{\mathrm{pose}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}},$$4. In erasing, vision peaks before contact, vision and force are jointly emphasized during wiping, and pose dominates during retreat (Feng et al., 19 May 2026). The paper interprets this as evidence that SO-TA learns the expected phase structure of contact-rich manipulation.

6. Broader OT-attention context, conceptual relatives, and limitations

SO-TA belongs to a wider body of work that treats attention as constrained transport rather than only softmax-normalized similarity. In textual network embedding, OT has been used as a global, sparse, self-normalized mutual attention mechanism, with the transport plan aligning token sequences under explicit marginal constraints (Chen et al., 2019). In slot-based object-centric learning, slot attention has been reinterpreted as an approximate OT solver, and cost shaping has been used to obtain low-entropy, tie-broken assignments in videos where the number of objects changes over time (Zhang et al., 2023). In training-free text-to-image synthesis, OT has been used to reposition cross-attention maps during denoising through a Sinkhorn solve over spatial distributions, specifically to mitigate mislocated objects (Han et al., 28 Mar 2025).

A distinct but closely related line addresses spatio-temporal measures rather than neural patch fusion. “Spatio-Temporal Alignments” defines a differentiable formulation in which soft-DTW handles temporal alignment and unbalanced OT measures spatial discrepancy, while a later barycentric formulation combines Soft-DTW with debiased UOT to compute spatio-temporal Fréchet means and explicitly address shifts in time, space, and total population size (Janati et al., 2019, Janati et al., 2022). These methods are conceptually close to a spacetime OT-attention interpretation, but they are not tri-modal visuo-haptic fusion backbones.

From a systems perspective, “FlashSinkhorn” shows that stabilized log-domain Sinkhorn half-steps for squared-Euclidean OT are exactly row-wise LogSumExp normalizations of biased dot-product scores, the same normalization as transformer attention. This allows FlashAttention-style fusion and tiling, with fused Triton kernels that reduce HBM traffic while retaining linear-memory operations (Ye et al., 3 Feb 2026). A plausible implication is that if SO-TA-style modules can preserve a row-wise LogSumExp structure or another streaming-friendly reduction, comparable fused-kernel benefits may become available.

The limitations reported for SO-TA are specific. The OT-derived heatmaps can remain somewhat diffuse because the cost $$\mathbf{Z}^{\mathrm{force}},\mathbf{Z}^{\mathrm{pose}}\in\mathbb{R}^{B\times T_w\times d_{\text{model}}},$$5 lacks any explicit spatial prior beyond cosine similarity (Feng et al., 19 May 2026). The structured fusion improves robustness, but it does not completely eliminate compounding-error tails under contact; the cutoff-and-reset wrapper closes the long tail at runtime, but this is an external safeguard rather than a learned solution (Feng et al., 19 May 2026). More generally, the paper’s evidence supports SO-TA as a principled multimodal inductive bias for contact-rich manipulation, not as a claim that OT alone resolves all sources of partial observability, distribution shift, or long-horizon control error.