Technical Summary: Rotary Masked Autoencoders are Versatile Learners

Problem Statement

The application of Transformer architectures to irregular time-series data presents significant challenges. Standard Transformers, originally designed for discrete sequences like natural language, rely on quantized positional information. They lack native support for continuous positional data, which is essential for modeling irregularly sampled time-series found in domains such as health, finance, and astrophysics. Existing solutions typically require specialized architectural modifications (e.g., altering feedforward layers) or novel positional embeddings (e.g., neural ODEs), which increase computational overhead and method complexity. Furthermore, extending Transformers to irregular data often necessitates abandoning the standard "ecosystem" of tools developed for Natural Language Processing (NLP) and Computer Vision (CV).

Methodology: Rotary Masked Autoencoder (RoMAE)

The authors propose RoMAE, an extension of the Masked Autoencoder (MAE) framework that integrates Rotary Positional Embeddings (RoPE) to handle continuous, multidimensional positional information without requiring task-specific architectural specializations.

Core Components

1. N-Dimensional Patchification: Inputs $x \in \mathbb{R}^{d_1 \times d_2 \times \dots \times d_D}$ are divided into non-overlapping patches. For irregular dimensions (e.g., time), the patch size must be 1 to ensure uniform data points per patch. These patches are flattened into a sequence and projected into an embedding dimension $d_{model}$.
2. Continuous Axial RoPE: Unlike standard RoPE designed for discrete integers, RoMAE utilizes RoPE with continuous real-valued positions ($m \in \mathbb{R}$). The model accepts a sequence of positional vectors $s = [s_1, \dots, s_k]$ alongside input values.
   • Axial Application: To encode multi-dimensional positions (e.g., time and channel index), the model employs Axial RoPE, splitting the embedding space into subspaces, each encoding a specific dimension.
   • Handling High Dimensions: To avoid the computational cost of dividing $d_{model}$ by a large number of dimensions $D$, the authors optionally reserve a dimension to store the dimensional index, reducing the effective positional dimensions.
3. Architecture: RoMAE follows the asymmetric encoder-decoder structure of MAE.
   • Encoder: A standard Transformer encoder (utilizing SiLU activation and RMSNorm) processes unmasked patches.
   • Decoder: A smaller Transformer decoder reconstructs masked patches.
   • Pre-training: A percentage of patches are masked. The decoder predicts the original patch values.
   • Downstream Tasks: The decoder is discarded, and the encoder's intermediate representations are used for classification or regression.

Theoretical Insights on Position

The paper investigates the interplay between RoPE (a Relative Positional Embedding method) and the [CLS] token:

• Relative vs. Absolute: Without a learned [CLS] token, the RoMAE encoder is translation-invariant, relying solely on relative positions.
• Reconstruction of Absolute Position: The authors prove (Proposition 4.2) that including a learned [CLS] token allows the model to reconstruct absolute positions. The [CLS] token acts as an "anchor," enabling the model to infer the absolute position of any embedding relative to the anchor, effectively breaking the pure relative property of RoPE in the context of the full sequence.

Key Contributions

1. Continuous Positional Embedding with RoPE: The work expands RoPE's utility to continuous, non-uniform timestamps, demonstrating that learned input embeddings can invalidate RoPE's relative distance guarantees unless an anchor (like [CLS]) is used.
2. RoMAE Framework: A unified architecture that natively handles irregular multivariate time-series while maintaining standard performance on images and audio. It eliminates the need for specialized time-series architectures.
3. Comprehensive Experimental Analysis: The authors validate RoMAE across four distinct modalities and tasks:
   • Irregularly sampled multivariate time-series classification.
   • Image classification.
   • Irregularly sampled time-series interpolation.
   • Audio classification.

Experimental Results

The paper reports results on several benchmarks, comparing RoMAE against state-of-the-art (SOTA) specialized models:

• Irregular Time-Series Classification (ELAsTiCC Challenge): On the DESC ELAsTiCC dataset (astronomical light curves), RoMAE-tiny outperformed the specialized ATAT architecture. RoMAE achieved an F-score of 0.8029 (without alert mask) compared to ATAT's 0.6270, demonstrating that a generalist model can surpass specialized ones on difficult datasets.
• Image Classification (Tiny ImageNet): RoMAE with RoPE performed comparably to or slightly better than MAE with absolute positional embeddings, confirming its versatility in CV.
• Audio Classification (ESC-50): Pre-trained on AudioSet-20k and Librispeech, RoMAE-small achieved 84.7% accuracy, outperforming the SSAST model (82.2%) trained on the same dataset.
• Interpolation Tasks:
  • Spirals: RoMAE improved upon the previous best (ContiFormer) by an order of magnitude in MSE.
  • Synthetic & PhysioNet: RoMAE achieved competitive MSE results on univariate and multivariate clinical data, often surpassing specialized models like HeTVAE and Latent ODE, particularly in handling sparse observations.
• Position Reconstruction: Experiments confirmed that models with a [CLS] token could reconstruct absolute positions with near-zero Mean Squared Error (MSE), while models without it failed (MSE > 200).

Significance and Claims

The paper claims that RoMAE is a versatile learner capable of bridging the gap between standard Transformer applications (NLP, CV) and irregular time-series analysis.

• Unified Framework: It demonstrates that specialized architectures are not strictly necessary for strong performance on irregular time-series; standard Transformer methods, augmented with continuous RoPE, suffice.
• Data Efficiency: RoMAE shows strong performance even on small datasets (hundreds of samples) and scales to millions (ELAsTiCC), highlighting the benefits of pre-training.
• Theoretical Insight: The work provides empirical evidence that the "long-term decay" property of RoPE is not critical for its function and clarifies the conditions under which RoPE can support absolute position reconstruction.
• Limitations: The authors acknowledge that continuous RoPE incurs a computational overhead (approx. 13% slower than quantized RoPE) because frequencies cannot be cached between forward passes. Additionally, like BERT-style models, RoMAE is not well-suited for extrapolation tasks involving out-of-distribution temporal windows.

In conclusion, RoMAE offers a task-agnostic baseline for continuous-time interpolation and representation learning, suggesting that the "foundation model" approach can be successfully extended to scientific fields dealing with irregular, multidimensional data without bespoke architectural changes.