Soft-MSM: Differentiable Context-Aware Elastic Alignment for Time Series

1. Problem Statement

Time series machine learning (TSML) relies heavily on distance measures to handle temporal misalignments between sequences. Elastic distances, such as Dynamic Time Warping (DTW) and Move-Split-Merge (MSM), address this by finding optimal alignment paths that stretch or compress time axes.

• The Limitation of Elastic Distances: While effective for tasks like clustering and classification, traditional elastic distances are non-differentiable. This prevents their direct use as loss functions in gradient-based optimization (e.g., neural networks, iterative averaging).
• The Limitation of Existing Solutions: Soft-DTW was introduced to make DTW differentiable by replacing the hard minimum in its recursion with a smooth soft-minimum operator. However, Soft-DTW cannot be directly extended to other elastic distances like MSM. MSM is often superior to DTW in performance because it uses context-aware transition costs (split/merge penalties) that depend on the local ordering of three values. The piecewise nature of MSM's cost function creates non-smooth regions that standard soft-minimum approaches cannot easily differentiate through.

2. Methodology: Soft-MSM

The authors propose Soft-MSM, a smooth, differentiable relaxation of the Move-Split-Merge distance. The core challenge was to make the MSM transition cost differentiable while preserving its context-aware structure.

Key Technical Innovations:

1. Smooth Gated Surrogate for Transition Costs:

   • The original MSM cost function applies a fixed penalty $c$ if a point $x$ lies between neighbors $y$ and $z$, and an additional penalty otherwise. This is a piecewise function with non-differentiable boundaries.
   • Solution: The authors introduce a smooth between-gate function $g(x, y, z)$ using a hyperbolic tangent-like approximation:
     $$g(x, y, z) = \frac{1}{2} \left( 1 - \frac{u}{\sqrt{u^2 + \epsilon}} \right), \quad \text{where } u = (x-y)(x-z)$$
   • This gate smoothly transitions between 1 (when $x$ is between $y$ and $z$) and 0 (otherwise), allowing gradients to flow through the decision boundary.
   • The transition cost is redefined as:
     $$\text{trans}_\gamma(x, y, z; c) = c + (1 - g(x, y, z)) \cdot \text{softmin}_\gamma((x-y)^2, (x-z)^2)$$
   • They also replace the absolute difference in the local distance with a squared difference to ensure differentiability at $x=y$.

2. Forward and Backward Recursions:

   • Forward Pass: Computes the accumulated cost matrix $C$ using a soft-minimum operator ($\text{softmin}_\gamma$) instead of the hard minimum, integrating the new smooth transition costs.
   • Backward Pass: Derives the soft alignment matrix $A$ and the Jacobian (gradient) with respect to the input time series. The gradient is decomposed into contributions from diagonal (match), vertical (split), and horizontal (merge) edges, weighted by the alignment probabilities.

3. Theoretical Properties:

   • Differentiability: Proven to be continuously differentiable for $\gamma > 0$ and $\epsilon > 0$.
   • Limiting Behavior: As $\gamma \to 0$ and $\epsilon \to 0$, Soft-MSM converges pointwise to the standard MSM recursion (with squared deviations).
   • Metricity: Soft-MSM is not a metric (it violates the identity of indiscernibles and can be negative). To address this, the authors propose a divergence-corrected formulation (analogous to Soft-DTW divergence) to ensure non-negativity and zero self-distance.
   • Complexity: Maintains the same $O(m^2)$ time and space complexity as standard MSM.

3. Key Contributions

• First Differentiable MSM: The first formulation of a differentiable version of the Move-Split-Merge distance, enabling its use in gradient-based learning.
• Novel Smooth Gate: A mathematical derivation of a smooth gating function that handles the context-dependent split/merge penalties of MSM without losing the structural benefits of the original distance.
• Closed-Form Gradients: Derivation of the Jacobian transformation allowing direct optimization of time series sequences using Soft-MSM.
• Implementation: The method is implemented in the open-source aeon toolkit, supporting reproducibility.

4. Experimental Results

The authors evaluated Soft-MSM on 112 UCR time series datasets across three main scenarios:

A. Time Series Averaging (Barycentre)

• Goal: Find a representative time series (barycentre) that minimizes the sum of distances to a set of input series.
• Comparison: Soft-MSM Barycentre (Soft-MBA) vs. standard MSM Barycentre (MBA) and Stochastic Subgradient MBA (SSG-MBA).
• Result: Soft-MBA achieved lower MSM Fréchet loss than both MBA and SSG-MBA on 97.2% to 100% of datasets (depending on the smoothing parameter $\gamma$). It significantly outperformed the hard counterparts, demonstrating that the differentiable relaxation leads to better convergence and lower loss.

B. Time Series Clustering

• Goal: Use k-means with elastic centroids.
• Comparison: Soft-MBA vs. Soft-DTW Barycentre (Soft-DBA), k-Shape, and Euclidean k-means.
• Result:
  • Soft-MBA significantly outperformed Soft-DBA in both Clustering Accuracy and Adjusted Rand Index (ARI).
  • Soft-MBA was competitive with k-Shape (a strong baseline), despite using a fundamentally different similarity measure.
  • Soft-MBA showed greater robustness to the smoothing parameter $\gamma$ compared to Soft-DBA.

C. Nearest-Centroid Classification

• Goal: Classify time series based on distance to class prototypes (barycentres).
• Result: Classifiers using Soft-MBA prototypes achieved significantly higher accuracy and balanced accuracy than those using Soft-DBA prototypes.
• Qualitative Analysis: On the CRICKETX dataset, Soft-MBA produced smoother, more interpretable class prototypes (capturing the "six runs" signal clearly) compared to the noisy prototypes generated by other methods.

5. Significance and Impact

• Bridging the Gap: Soft-MSM successfully bridges the gap between high-performance elastic distances (like MSM) and modern gradient-based deep learning frameworks.
• Superiority over Soft-DTW: The results suggest that for many time series tasks, MSM is a superior distance measure to DTW, and making it differentiable yields better results than simply using Soft-DTW.
• Broad Applicability: The method opens the door for using MSM-based losses in:
  • Deep time series forecasting.
  • Generative models for data augmentation.
  • Weakly supervised alignment and segmentation.
  • Representation learning (e.g., Series2Vec).
• Inductive Bias: It provides a new inductive bias based on "move, split, and merge" operations, which is particularly useful for piecewise-constant or event-driven time series data.

In conclusion, Soft-MSM is a significant advancement in time series machine learning, offering a differentiable, context-aware elastic distance that outperforms existing Soft-DTW-based approaches in averaging, clustering, and classification tasks.