# ReLATE: Learning Efficient Sparse Encoding for High-Performance Tensor Decomposition

## TL;DR

ReLATE is a reinforcement learning framework that automatically learns efficient sparse tensor encodings, significantly improving tensor decomposition performance on irregular, high-dimensional sparse data without requiring labeled training data.

## Contribution

It introduces a novel reinforcement learning-based approach for automatically constructing optimized sparse tensor formats tailored to data and computation needs.

## Key findings

- Achieves up to 2X speedup over expert-designed formats.
- Automatically adapts to irregular tensor shapes and data distributions.
- Provides consistent performance improvements across diverse datasets.

## Abstract

Tensor decomposition (TD) is essential for analyzing high-dimensional sparse data, yet its irregular computations and memory-access patterns pose major performance challenges on modern parallel processors. Prior works rely on expert-designed sparse tensor formats that fail to adapt to irregular tensor shapes and/or highly variable data distributions. We present the reinforcement-learned adaptive tensor encoding (ReLATE) framework, a novel learning-augmented method that automatically constructs efficient sparse tensor representations without labeled training samples. ReLATE employs an autonomous agent that discovers optimized tensor encodings through direct interaction with the TD environment, leveraging a hybrid model-free and model-based algorithm to learn from both real and imagined actions. Moreover, ReLATE introduces rule-driven action masking and dynamics-informed action filtering mechanisms that ensure functionally correct tensor encoding with bounded execution time, even during early learning stages. By automatically adapting to both irregular tensor shapes and data distributions, ReLATE generates sparse tensor representations that consistently outperform expert-designed formats across diverse sparse tensor data sets, achieving up to 2X speedup compared to the best sparse format, with a geometric-mean speedup of 1.4-1.46X.