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.

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.

Figures

Tables

Table 1. Table 1 . Characteristics of the target sparse tensor data sets.

Tensor	Dimensions	#NNZs	Density
darpa	$22.5K\times 22.5K\times 23.8M$	$28.4M$	$2.4\times {10}^{-09}$
fb-m	$23.3M\times 23.3M\times 166$	$99.6M$	$1.1\times {10}^{-09}$
flickr-4d	$319.7K\times 28.2M\times 1.6M\times 731$	$112.9M$	$1.1\times {10}^{-14}$
flickr-3d	$319.7K\times 28.2M\times 1.6M$	$112.9M$	$7.8\times {10}^{-12}$
deli-4d	$532.9K\times 17.3M\times 2.5M\times 1.4K$	$140.1M$	$4.3\times {10}^{-15}$
deli-3d	$532.9K\times 17.3M\times 2.5M$	$140.1M$	$6.1\times {10}^{-12}$
nell-1	$2.9M\times 2.1M\times 25.5M$	$143.6M$	$9.1\times {10}^{-13}$
amazon	$4.8M\times 1.8M\times 1.8M$	$1.7B$	$1.1\times {10}^{-10}$
patents	$46\times 239.2K\times 239.2K$	$3.6B$	$1.4\times {10}^{-03}$
reddit	$8.2M\times 177K\times 8.1M$	$4.7B$	$4.0\times {10}^{-10}$

Table 2. Table 2 . ReLATE Hyperparameters.

Hyperparameter	Value	Description
$\gamma$	$0.99$	Discount factor for future rewards.
$\left|ℬ\right|$	$\ell \left(p\right)$	Number of samples per minibatch.
$\left|{ℬ}^{\mathrm{PER}}\right|$	$1M$	Max. number of observations in the replay buffer.
Initial $\alpha$	$0.001$	Initial learning rate before decay.
Min. $\alpha$	$0.0001$	Learning rate after decay.
Initial $ϵ$	$1.0$	Initial value for $ϵ$ before decay.
Min. $ϵ$	$0.1$	Value of $ϵ$ after decay.
${\tau }_{\mathrm{update}}$	$100$	Target update frequency.
$E_{max}$	$5000$	Number of training episodes.
Min. accuracy	$90\%$	Accuracy threshold for the reward model.
Optimizer	Adam	Training optimizer used for model updates.

Table 3. Table 3 . Comparison of memory metrics (obtained by perf counters) across ReLATE and ALTO on a 128-core EMR system.

Tensor	Format	Memory Volume	L2 Miss Ratio	L3 Miss Ratio
deli-3d	ReLATE	$53.34\text{GB}$	$68.56\%$	$59.95\%$
	ALTO	$93.69\text{GB}$	$75.03\%$	$78.82\%$
patents	ReLATE	$176.67\text{GB}$	$52.05\%$	$72.96\%$
	ALTO	$176.74\text{GB}$	$55.70\%$	$89.41\%$
reddit	ReLATE	$437.3\text{GB}$	$44.03\%$	$33.39\%$
	ALTO	$743.6\text{GB}$	$49.50\%$	$43.70\%$

Equations

$$\left(\genfrac{}{}{0ex}{}{\ell (p)}{\ell (1),\ell (2),\dots ,\ell (N)}\right)=\frac{\ell (p)!}{\ell (1)!\ell (2)!\cdots \ell (N)!}$$

$$a_t={\mathrm{argmax}}_{a_t\in A^{-}}Q(s_t,a_t;\theta )$$

$$\frac{1}{|\mathcal{B}|}\sum_{i=1}^{|\mathcal{B}|}\Bigl{[}r_{i}+\gamma\,Q\!\bigl{(}s_{i}^{\prime},\operatorname{argmax}_{a^{\prime}}Q(s_{i}^{\prime},a^{\prime};\theta);\theta^{-}\bigr{)}-Q(s_{i},a_{i};\theta)\Bigr{]}^{2}$$

Full text

Sparse tensors, tensor decomposition, MTTKRP, multi-core CPU, reinforcement learning, learned format

††conference: ; ††journalyear: ;††copyright: none

1. Introduction

Tensor decomposition (TD) generalizes principal component analysis to break down complex, high-dimensional sparse data, such as financial transactions, electronic health records, and user ratings (Kobayashi et al., 2018; Yadav et al., 2018; Symeonidis, 2016), into simpler, low-dimensional components, which reveal the latent relationships and trends within the data. Although TD algorithms play a critical role in dimensionality reduction, compression, and analysis of multi-way data (Kolda and Bader, 2009; Sidiropoulos et al., 2017; Panagakis et al., 2021; Liu and Parhi, 2023), they remain challenging to scale on modern parallel processors because of their highly irregular computations and memory-access patterns (Choi et al., 2018; Helal et al., 2021). Additionally, sparse tensors can be represented in a multitude of ways, each tailored for different sparsity patterns and exhibiting varying storage and performance trade-offs (Smith and Karypis, 2017; Laukemann et al., 2025). Moreover, TD algorithms perform computations across all dimensions (i.e., modes) of a sparse tensor, and even if a sparse representation is efficient for a particular mode, it might deliver subpar performance in other modes (Smith and Karypis, 2017; Helal et al., 2021).

Since TD problems are both high-dimensional and data-dependent, their optimization space is vast ($\S$3.3.2). However, the state-of-the-art approaches rely on suboptimal heuristics for data representation and parallel execution that fail to account for the irregular shapes and/or data distributions of sparse tensors. Prior work improved the parallel performance of TD using compressed sparse tensor representations, including the mode-specific CSF format (Smith et al., 2015; Smith and Karypis, 2017) and the mode-agnostic HiCOO format (Li et al., 2018, 2019). While these sparse representations realized substantial performance gains compared to the de facto COO format, they are oblivious to the highly irregular shapes and data distributions of sparse tensors (Helal et al., 2021; Laukemann et al., 2025). The state-of-the-art linearized tensor formats (Helal et al., 2021; Nguyen et al., 2022; Soh et al., 2023; Laukemann et al., 2025), such as the mode-agnostic ALTO format, tailor their tensor representation to the irregular shapes of sparse tensors, yet they overlook the underlying data distributions. SpTFS (Sun et al., 2020, 2021) leverages supervised machine learning methods to predict the best of COO, HiCOO, and CSF formats to compute the decomposition of a given sparse tensor. However, due to the lack of large-scale sparse tensor training sets and because the optimal solution (or ground truth) for such irregular workloads is unknown, supervised learning methods are impractical. Additionally, the performance of these methods is bounded by the best existing formats, even if they managed to attain oracle-level prediction accuracy.

In place of suboptimal, input-agnostic heuristics, we propose the reinforcement-learned adaptive tensor encoding (ReLATE) framework. ReLATE adopts the emerging learned or learning-augmented paradigm (Mitzenmacher and Vassilvitskii, 2022), in which we devise a skeleton algorithm that integrates domain expertise through a rich problem formulation to dynamically navigate the solution space in search of an effective solution. In contrast to classical auto-tuning techniques, our learned algorithm goes beyond adjusting program parameters to construct an efficient sparse tensor encoding along with the corresponding schedule of tensor compute and memory operations.

Central to ReLATE is an autonomous agent that leverages deep reinforcement learning (DRL) and domain knowledge to improve the performance and scalability of high-dimensional data analytics without requiring labeled training data sets. Specifically, our ReLATE agent directly learns the best sparse tensor encoding from offline interactions with the TD environment via actions and from observing the environment state and the reward signal. Most importantly, ReLATE introduces a hybrid model-free and model-based learning algorithm, where an agent employs an adaptive neural network architecture that scales with environment complexity and learns from both real (evaluated) actions as well as imagined (estimated) actions to effectively handle large-scale tensors. Additionally, ReLATE accelerates the learning process using rule-driven action masking and dynamics-informed action filtering mechanisms to prune invalid actions during the early exploration stage and low-value actions in the delayed exploitation stage of the learning process. For accurate reward measurements in high-performance computing (HPC) environments that are both growing in complexity and prone to noise (Tuncer et al., 2017), ReLATE uses a relative reward signal and adopts a decoupled server-client execution model, where the learning and decision making is isolated from the reward evaluation by running the client and server on separate compute nodes.

By overcoming the limitations of prior supervised learning methods and tensor formats, ReLATE automatically adapts to the irregular shapes and data distributions of sparse tensors and to the target hardware architecture. As a result, ReLATE discovers sparse representations that outperform the expert-designed formats across a representative set of sparse tensors. In summary, the main contributions of this work are:

• •

We introduce ReLATE, a novel DRL framework for constructing efficient mode-agnostic sparse tensor representations without requiring any labeled examples. In contrast to prior work (Smith et al., 2015; Smith and Karypis, 2017; Li et al., 2018; Helal et al., 2021; Laukemann et al., 2025), ReLATE automatically adapts not only to the asymmetric shapes of sparse tensors but also to their highly irregular data distributions ($\S$3).

• •

We tackle the exploration-exploitation dilemma by designing a domain-specific agent that can effectively navigate the vast solution space using rule-driven action masking and model-guided action filtering mechanisms. Hence, ReLATE always generates functionally correct sparse representations, with bounded execution time, even before learning the environment dynamics, enabling deployment in live sparse tensor environments ($\S$3 and $\S$4).

• •

Across a diverse set of real-world sparse tensors, ReLATE outperforms the prior state-of-the-art formats for both the original as well as randomly permuted data sets, delivering $1.4\times$ and $1.46\times$ geometric-mean speedup, respectively, over the best format on an Intel Emerald Rapids system. Additionally, our learned encoding realizes up to $2\times$ speedup over the best format for large-scale, low-density tensors, which are particularly challenging to optimize on modern parallel processors ($\S$4).

2. Background and Related Work

2.1. Tensor Notations

A tensor is an $N$-dimensional array, with each element addressed by an $N$-tuple index $\textbf{i}~=~(i_{1},i_{2},\ldots,i_{N})$. The coordinate $i_{n}$ locates a tensor element along the $n^{\mathrm{th}}$ mode, whose length is $I_{n}$, with $n\in\{1,2,\ldots,N\}$ and $i_{n}\in\{0,1,\ldots,I_{n}-1\}$. A tensor is considered sparse when most of its elements are zero. The following notations are used in this paper:

1. (1)

Tensors are denoted by Euler script letters (e.g. $\mathcal{X}$). 2. (2)

Fibers in a tensor generalize matrix rows and columns. A mode-$n$ fiber of a tensor $\mathcal{X}$ is any vector obtained by fixing all indices of $\mathcal{X}$ except the $n^{th}$ index. 3. (3)

Tensor matricization is the process of unfolding a tensor into a matrix. The mode-$n$ matricization of a tensor is denoted as $\mathbf{X}_{(n)}$, and is obtained by arranging all the mode-$n$ fibers of $\mathcal{X}$ as the columns of $\mathbf{X}_{(n)}$. 4. (4)

The Khatri-Rao product (KRP) (Liu and Trenkler, 2008) is the column-wise Kronecker product of two matrices, and is denoted by $\odot$. Given matrices $\mathbf{C}^{(1)}$ $\in$ $\mathbb{R}^{I_{1}\times F}$ and $\mathbf{C}^{(2)}$ $\in$ $\mathbb{R}^{I_{2}\times F}$, their KRP is $\mathbf{K}=\mathbf{C}^{(1)}\odot\mathbf{C}^{(2)}=\left[\textbf{c}_{1}^{(1)}\otimes\textbf{c}_{1}^{(2)}\ \textbf{c}_{2}^{(1)}\otimes\textbf{c}_{2}^{(2)}\dots\textbf{c}_{F}^{(1)}\otimes\textbf{c}_{F}^{(2)}\right]$, where $\otimes$ denotes the Kronecker product, $\mathbf{c}_{f}$ denotes the $f^{th}$ factor or column of the matrix $\mathbf{C}$, and $\mathbf{K}$ $\in$ $\mathbb{R}^{(I_{1}\cdot I_{2})\times F}$.

2.2. Tensor Decomposition (TD)

Tensor decomposition (TD) is a generalized form of singular value matrix decomposition and principal component analysis. This work targets algorithms that iteratively compute the most popular TD model for sparse tensors, the canonical polyadic decomposition (CPD), which decomposes a tensor into a finite sum of rank‑one tensors, each representing a principal component. The survey by Kolda and Bader (Kolda and Bader, 2009) provides a thorough overview of TD algorithms.

The key performance bottleneck in the target TD algorithms is the matricized tensor times Khatri-Rao product (MTTKRP) operation (Smith et al., 2015; Laukemann et al., 2025), which must be executed along each mode in every iteration. The mode-$n$ MTTKRP on a tensor $\mathcal{X}$ is defined as: $\mathbf{X}_{(n)}\bigodot_{k=1,k\neq n}^{N}\mathbf{C}^{(k)}$, where $\mathbf{C}^{(n)}\in$ $\mathbb{R}^{I_{n}\times F}$ $\forall n\in\{1,2,\ldots,N\}$ are the factor matrices for each mode.

Previous research focused on accelerating the parallel execution of MTTKRP using various sparse tensor formats, which can be classified into mode-specific and mode-agnostic representations. Mode-specific formats (Smith et al., 2015; Smith and Karypis, 2015, 2017; Nisa et al., 2019b, a; Kurt et al., 2022) generalize the classical compressed sparse row (CSR) format and typically require multiple tensor copies for best performance. Mode-agnostic formats (Li et al., 2018, 2019; Helal et al., 2021; Nguyen et al., 2022; Soh et al., 2023; Laukemann et al., 2025) maintain a single tensor copy by storing nonzero elements and their compressed coordinates in list-based representations. Laukemann et al. (Laukemann et al., 2025) provide a detailed review of sparse tensor formats and benchmark their performance on TD kernels.

2.3. Deep Reinforcement Learning (DRL)

Reinforcement learning (RL) is a goal-directed, decision-making process that emerged in biological brains to enable flexible responses to their environment (Bennett, 2023). DRL combines the trial‐and‐error learning method of RL with the automatic feature selection and hierarchical representation learning of modern deep neural networks (Mnih et al., 2013, 2015). In contrast to common misperception, trial-and-error learning is not random, as information about the target environment grows with every trial, especially when guided by a model capable of retaining favorable outcomes and discarding unfavorable ones. Specifically, a DRL-based agent learns a data-driven policy by interacting with the environment through actions and observing the environment state and the corresponding reward or punishment signal. Formally, this Markov Decision Process (MDP) is defined by a tuple $\langle S,A,R,T,\gamma\rangle$, where $S$ is a set of environment states, $A$ is a set of actions, $R(s_{t},a_{t})$ is a reward function, $T(s_{t},a_{t},s_{t+1})$ is a transition function, and $\gamma\in[0,1]$ is a discount factor. At each step or time instance $t$, the agent moves the environment from its current state $s_{t}\in S$ to a new state $s_{t+1}\in S$ by selecting an action $a_{t}\in A$, and then it receives a reward $r_{t}=R(s_{t},a_{t})$. Hence, an interaction (or observation) is defined by a tuple $\langle s_{t},a_{t},r_{t},s_{t+1}\rangle$. A policy $\pi$ determines which action to take for each environment state $s_{t}\in S$. The goal of the agent is to learn the policy $\pi^{*}$ that maximizes the cumulative future reward obtained from the environment, that is, the sum of rewards $r_{t}$ discounted by a factor $\gamma$ at each time step $t$.

One issue in this learning approach is sample inefficiency, i.e., the agent typically requires a large number of interactions with its environment to learn an effective policy. To improve sample efficiency, deep Q-network (DQN) (Mnih et al., 2013, 2015) introduces an experience replay buffer to collect the most recent observations and then learns by sampling minibatches from the replay buffer instead of learning only from the current experience. Additionally, this replay mechanism enables DQN to learn effective policies in an offline setting with limited (or user-defined) memory capacity. Specifically, if the value $Q(s_{t},a_{t})$ is the future reward of a state-action pair, the optimal policy $\pi^{*}=\operatorname{argmax}_{a_{t}\in A}Q^{*}(s_{t},a_{t})$ selects the highest valued action in all states. DQN approximates the value function $Q(s_{t},a_{t})$ using a deep neural policy network that estimates the values of all actions in a given state. To improve learning stability, double DQN (Hasselt et al., 2016) uses a policy network to select the best action and employs a target network, which is updated periodically, to estimate the future reward of this action. Instead of uniform random sampling from the experience replay buffer, the prioritized replay mechanism (Schaul et al., 2016) samples important experiences more frequently to further improve and accelerate the learning process (Hessel et al., 2018).

DRL has recently found applications in many performance-centric problems, notably in code generation (Mankowitz et al., 2023; Fawzi et al., 2022), compiler optimization (Mammadli et al., 2020; Ahn et al., 2019), and job scheduling (Gu et al., 2025; Mangalampalli et al., 2024). In the domain of sparse matrix and tensor formats, researchers have applied DRL to automate format selection and tuning (Armstrong and Rendell, 2008; Chen et al., 2023), yet the majority of prior works formulate these tasks as classification or regression problems and leverage supervised learning methods (Sun et al., 2020, 2021; Zhao et al., 2018; Xie et al., 2019; Stylianou and Weiland, 2023; Xiao et al., 2024; Yesil et al., 2022; Gao et al., 2024).

3. ReLATE Framework

3.1. Motivation and Challenges

Devising a learnable sparse tensor format is challenging because of the vast design space of this high-dimensional, data-dependent problem as well as the overhead and complexity involved in computing the reward. Accordingly, it can be computationally prohibitive for DRL agents to explore such a massive state-action space and to learn effective policies from a noisy, high-latency reward signal (Mnih et al., 2015).

Moreover, using existing sparse formats as a baseline for learned sparse tensor encoding is infeasible since they result in state–action spaces that are either too small or excessively large. For instance, the mode-specific CSF format (Smith et al., 2015; Smith and Karypis, 2017) splits the tensor into slices and fibers along a given mode order, and thus its performance largely depends on which mode is selected as the root mode. Hence, a CSF-based format exposes a limited number of options for the learner model. While block-based formats (Li et al., 2018, 2019) are mode-agnostic, setting the block size as the learnable parameter leads to a large state-action space proportional to the length of the tensor modes. Furthermore, unlike sparse matrices, the block size typically has limited performance impact on sparse tensors, which exhibit unstructured and extreme sparsity because of the curse of dimensionality (Helal et al., 2021; Laukemann et al., 2025). Similarly, the number of ways to represent a sparse tensor using linearized tensor formats (Helal et al., 2021; Nguyen et al., 2022; Laukemann et al., 2025) not only grows with the length of the tensor modes, but also with the number of modes.

Even with a tractable state–action space, evaluating the reward for every interaction with the TD environment incurs significant difficulty, due to the cost of executing TD kernels on the target hardware architecture. Additionally, modern HPC systems suffer from high run-time variability, driven by hardware complexity, resource contention, and operating system jitter (Tuncer et al., 2017), which can significantly affect the reward accuracy and thus the convergence of DRL agents.

3.2. ReLATE Overview

To overcome the challenges of learning how to efficiently represent high-dimensional sparse data, we introduce the ReLATE framework. Figure 1 presents the high-level view of the proposed DRL-based approach for encoding sparse tensors, which leverages an off-policy learning algorithm based on stored experience and adopts a decoupled server-client execution. By separating sample collection from the policy being learned, ReLATE enables offline learning, independently of the agent’s actions, and improves sample efficiency. When the client (agent) and server (environment) are deployed on different compute nodes, ReLATE isolates learning and decision making from reward evaluation to reduce environment noise and to improve reward accuracy. In addition, this decoupled execution allows ReLATE to utilize idle data center capacity to improve the performance of TD workloads without disrupting ongoing data analysis jobs.

Specifically, sparse tensor encoding, along with the associated TD algorithm determine the schedule of compute and memory operations, and in turn the parallel performance. The proposed learned encoding leverages the state-of-the-art ALTO format (Helal et al., 2021; Laukemann et al., 2025), due to its mode-agnostic nature and superior performance compared to prior formats. However, linearized tensor formats have a massive state-action space that is exponential in the number of tensor modes and mode lengths, which is intractable for DRL agents. Thus, ReLATE devises an environment representation that reduces the state-action space without severely restricting the agent’s degrees of freedom. Furthermore, it models the problem of finding the best encoding of a sparse tensor as a sequential decision problem, where the DRL agent transitions the environment from its initial state to a terminal state through a series of actions. In this model, the reward is computed based on the terminal state, which represents an encoding of the target sparse tensor, as the speedup compared to ALTO. By interacting with the TD environment and collecting complete trajectories (i.e., sequences of observations ending with terminal states) along with their rewards, the agent learns the best performance-centric policy that maximizes the overall reward using multiple feedforward neural networks.

In this problem formulation, there are two bottlenecks: the training cost and reward evaluation overhead. The training cost is proportional to the state-action space, which depends on the tensor shape, i.e., the number of dimensions and the dimension lengths. The overhead of computing the reward is not only a function of the tensor shape but also the number of nonzero elements. To tackle complex, high-dimensional environments with extensive dimension lengths, ReLATE introduces novel model features and enhancements. Instead of using negative rewards (or punishments), ReLATE accelerates the learning process by using an action masking mechanism to ensure that every episode or trajectory ends with a valid tensor encoding. Thus, the agent learns only from valid episodes and after receiving a normalized reward at the end of every episode. In this daunting sparse/delayed‐reward environment, the normalized (or differential) reward signal not only improves the learning stability but also allows the agent to discriminate the promising trajectories from the multitude of possible ones. To achieve convergence, credit is computed from the delayed reward and assigned back to every action taken by the agent using the reward shaping function. Additionally, ReLATE employs convolutional neural networks (CNNs) to encode the hierarchical spatial information of the target environment and uses an adaptive neural network architecture, where the number of hidden units scales with the size of the state-action space.

Due to the substantial overhead of reward evaluation, ReLATE keeps a reward cache to prevent redundant evaluation of known valid encodings. Most importantly, it presents a novel action selection and filtering algorithm that allows early exploration of new actions and delayed exploitation of the knowledge embedded in the agent. That is, ReLATE uses a hybrid model-free and model-based reinforcement learning approach, in which the agent progressively builds a reward model of the environment based on its experience during the exploration stage. Once the accuracy of the reward model is high (e.g., more than 90%), the agent switches to learning from both real actions (reward evaluation) and imagined actions (reward model). Hence, if an imagined action is deemed high value, in terms of its potential reward, the agent evaluates the actual reward and uses it to further refine the reward model. This hybrid learning approach enables safe and fast navigation of the solution space by reducing the real slowdown (compared to the expert format) encountered during the exploration stage and by significantly accelerating the reward evaluation during the exploitation stage.

3.3. Problem Formulation

As discussed in $\S$3.1, prior sparse tensor formats are ill-suited for goal-directed learning as they are restricted to structured sparsity and/or their state–action spaces are either overly constrained or unmanageably large. ReLATE learns efficient sparse tensor representations by reformulating this hard problem into a tractable one, mapping the construction of a linearized tensor format into a Markov Decision Process (MDP) amenable to DRL-based methods.

The state-of-the-art linearized formats (Nguyen et al., 2022; Laukemann et al., 2025) impose a total locality-preserving, mode-agnostic order on the nonzero elements of a sparse tensor to enable compact storage coupled with fast tensor access and traversal. Formally, a linearized sparse tensor $\mathcal{X}$ =\{$${x}_{1}, ${x}_{2}$, $\dots$, {x}_{M}$$\} is an ordered collection of nonzero elements, where each element ${x}_{i}$ $=\langle v_{i},p_{i}\rangle$ has a numerical value $v_{i}$ and a position $p_{i}$ (Laukemann et al., 2025). The position $p_{i}$ of the $i^{th}$ nonzero element is a compact linear encoding of the N-tuple index, or coordinates, $\textbf{i}~=(i_{1},i_{2},\dots,i_{N})$, with $i_{n}\in\{0,1,\ldots,I_{n}-1\}$ and $n\in\{1,2,\ldots,N\}$. Precisely, the linear encoding $p$ uses $\ell(p)=\sum_{n=1}^{N}\ell(n)\text{ bits}$, with $\ell(n)=\lceil\log_{2}I_{n}\rceil$, to represent the coordinates of nonzero elements as a sequence of bits $(b^{\ell(p)-1}\dots b^{1}~b^{0})_{2}$, where $b^{j}$ denotes the $(j+1)^{th}$ bit of that encoding. The coordinates i of a nonzero element can be quickly derived from its position $p_{i}$, and vice versa, using simple bit-level operations that can be overlapped with memory-intensive tensor computations (Helal et al., 2021; Nguyen et al., 2022). In contrast, traditional space-filling curves (Peano, 1890), which typically target dense data, use $N\times\max_{n=1}^{N}\ell(n)\text{ bits}$ for linear indexing (Helal et al., 2021). This can be much more expensive, as real-world tensors generally have skewed shapes (see Table 1). The linearized ALTO format and its variants substantially reduce the indexing size using a non-fractal encoding scheme that adapts to the irregular shapes of sparse tensors. This encoding partitions the tensor space along the longest mode first, which is equivalent to interleaving bits from the N-dimensional coordinates, starting with the shortest mode (Laukemann et al., 2025). Figure 2 shows a $4\times 8\times 2$ tensor with six nonzero elements encoded in the COO and ALTO formats (Laukemann et al., 2025).

While the resulting one-dimensional layout of linearized formats accommodates the asymmetric shapes of sparse tensors, it ignores their nonuniform and highly skewed data distributions, which directly affect the data reuse, workload balance, and synchronization overhead of parallel TD operations (Helal et al., 2021; Laukemann et al., 2025). Hence, efficient high-dimensional data analysis requires a linearization scheme that matches the unique sparsity patterns of each tensor. However, given that, in general, each bit of the linear encoding $p$ can be selected to represent any bit of the N-tuple index i, the number of possible linearized representations of a tensor grows exponentially with its number of modes and mode lengths. Assuming, without loss of generality, that all tensor modes have the same length, there are $(N~\ell(n))!$ possible linear formats. For a small 4-dimensional tensor with a mode length that fits into a single byte, the number of possible linearized representations is astronomically large ($(4\times 8)!\approx 2.63\times 10^{35}$).

Since no single linearized format (or, more broadly, no compressed tensor format) is optimal across all sparse tensor data sets, ReLATE bridges this capability gap by searching a subset of the nearly infinite ways to linearly arrange the multi-dimensional data points of a tensor. Hence, instead of exploring $(N~\ell(n))!$ possible linearized formats, ReLATE considers (N~\ell(n))!\big{/}(\ell(n)!)^{N} interleaved linear encodings ($\approx 9.96\times 10^{16}$ for the prior example) to uncover effective sparse representations that significantly improve the performance of TD algorithms ($\S$3.3.2). More importantly, by sequentially decomposing the combinatorial choice of a linearized tensor encoding into $N~\ell(n)$ simple decisions, ReLATE transforms an intractable problem into an MDP, defined by environment states and actions as well as transition and reward functions, which can be tackled by DRL agents.

3.3.1. Environment States

ReLATE defines an encoding environment with three types of states (namely, initial, intermediate, and terminal state) to represent the target sparse tensor. The agent drives the environment from its initial state (no encoding) to a terminal state (a linearized encoding) by executing a sequence of actions. A state $s_{t}\in S$ is an encoding matrix with $N$ rows and $\ell(p)$ columns. Each column $j$ can have at most one hot bit, which indicates that the next low bit from a mode index $i_{n}$ (if the $n^{th}$ bit in column $j$ is hot) is mapped to the $(j+1)^{th}$ bit in the linear encoding (or bit $b^{j}$). Hence, in an intermediate state, the number of hot bits is strictly less than $\ell(p)$, while a terminal state has exactly $\ell(p)$ hot bits. Figure 3 illustrates this environment representation for the example sparse tensor from Figure 2, with a terminal state representing the ALTO-based linearized encoding. In this example, the terminal state indicates that the first and second bits from the mode index or coordinate $i_{1}$ (i.e., $b_{i_{1}}^{0}$ and $b_{i_{1}}^{1}$) are mapped to the second and fourth bits in the linear encoding $p$ (i.e., $b^{1}$ and $b^{3}$), respectively.

Although a more compact environment representation is possible using a vector of length $\ell(p)$, it requires integer encoding, instead of one-hot encoding, for tensor modes, which introduces a bias and leads to poor learning performance by implying an arbitrary ordinal relationship across modes where none exists. Additionally, generalizing the environment representation so that each cell $[n,j]$ in the encoding matrix indicates which bit from a mode index $i_{n}$ is mapped to the $(j+1)^{th}$ bit in the linear encoding $p$ leads to a prohibitively large ($\ell(p)!$) state space, as explained in $\S$3.3.2.

3.3.2. Actions

In every episode, the agent performs $\ell(p)$ actions to transition the environment from its initial (no encoding) state to a terminal state by selecting the bits of the linear encoding $p$. In a step $t\in\{0,1,\ldots,\ell(p)-1\}$, the agent picks a mode index $i_{n}$ out of the $N$ tensor modes to map its next low bit to the current bit of the linear index ($b^{t}$). Hence, the action space $A=\{a_{t}\in\mathbb{N}:1\leq a_{t}\leq N\}$ is discrete and limited to $N$ actions, corresponding to the number of tensor modes. By restricting the number of permitted actions, ReLATE decreases the estimation error in selecting the highest-valued action (Hasselt et al., 2016). More importantly, by always mapping the next low bit of a mode index and preserving the internal bit order within each mode index, ReLATE dramatically reduces the solution space of the possible linearized representations. Specifically, it transforms the problem of exploring all $\ell(p)!$ linearized formats to the task of interleaving $N$ bit sequences, with length $\ell(n)~\forall n\in\{1,2,\dots,N\}$, into a sequence of length $\ell(p)$, which is the classical multinomial problem (Stanley, 2011). Since the bits of every mode index $i_{n}$ represent a subsequence of the linear encoding $p$, the number of ways to generate this interleaved linear encoding is given by:

[TABLE]

Thus, for the toy example in Figure 3, ReLATE limits the number of actions to $3$ and reduces the state-action space from $720$ to $60$ linearized representations.

3.3.3. Transition Function

Once an action $a_{t}$ is selected, the transition function $T(s_{t},a_{t},s_{t+1})$ constructs a new environment encoding $s_{t+1}$ with $t+1$ hot bits. If $s_{t+1}$ is a terminal state with $\ell(p)$ hot bits, the transition function invokes the reward function to examine the new linear encoding.

3.3.4. Reward Function

To evaluate the effectiveness of a newly constructed linear tensor encoding, the reward function $R(s_{t},a_{t})$ sends this encoding to the TD environment and receives runtime information (i.e., the execution time), as shown in Figure 1. However, using the raw execution time for the reward leads to exploding/vanishing error gradients and unstable learning, even after extensive hyper-parameter tuning. Instead of clipping the reward into a limited range (Mnih et al., 2015), and sacrificing its fidelity, ReLATE retains the magnitude information by evaluating the new linear encoding against an expert policy. Hence, the reward function $R(s_{t},a_{t})$ generates a high-fidelity, differential reward signal as the speedup compared to the ALTO-based encoding (Laukemann et al., 2025), thereby resulting in bounded error gradients and allowing ReLATE to use the same hyper-parameters across diverse sparse tensors.

3.4. Learning Algorithm

Even though the proposed MDP formulation of sparse tensor workloads allows ReLATE to learn the best policy to minimize their execution time, it is still challenging to tackle complex environment characterized by a large number of nonzero elements, extensive mode length, and high dimensionality. The main reason is that the state-action space in such environments remains vast, even after substantial pruning, which in turn increases the training cost; moreover, each reward evaluation is expensive, making it essential to minimize environment interactions and to maximize the information extracted from every interaction. To address these challenges, ReLATE effectively combines model-free (double DQN (Hasselt et al., 2016)) and model-based reinforcement learning and introduces a set of new model features and optimizations to significantly accelerate the training process.

Algorithm 1, together with Figure 1, outlines the key concepts of the proposed learning strategy for generating optimized sparse tensor encoding. Given a sparse tensor $\mathcal{X}$ and a TD workload, ReLATE constructs an environment representation, based on the MDP formulation detailed in $\S$3.3, to find an optimized linear encoding $p^{*}$. ReLATE leverages a prioritized replay buffer (Schaul et al., 2016) to improve sample efficiency by keeping a diverse set of the latest significant transitions observed by the agent. To learn a stable encoding policy, ReLATE uses a policy network ($\theta$) along with a target network ($\theta^{-}$) to learn a mapping from the high-dimensional state space $S$ to the discrete action space $A$. As illustrated in Figure 4, the proposed DRL agent leverages an adaptive CNN architecture, in which the number of hidden units increases with the size of the state-action space, to encode the hierarchical spatial information of the sparse tensor environment. Both the policy and target networks comprise two convolutional layers, with $16$ and $32$ feature maps respectively, each using $3\times 3$ filters to capture the spatial correlations present within a receptive field of $5\times 5$, followed by two fully connected layers. The input layer of these networks ingests an $N\times\ell(p)$ encoding matrix representing the current environment state, while the output layer emits $N$ activations, each indicating the value of a possible action.

The training process unfolds over a number of episodes, each consisting of $\ell(p)$ steps. For every episode, ReLATE moves the environment from its initial state $s_{0}$ to a terminal state $s_{\ell(p)}$ and gathers transitions in a temporary memory buffer for further processing. In each step $t$, ReLATE ensures that the number of bits selected from a mode and mapped to the linear encoding (i.e., the number of hot bits in any row $n$ of the encoding matrix, as illustrated in Figure 3 and $\S$3.3) can not exceed the number of mode bits $\ell(n)~\forall~n\in\{1,2,\ldots,N\}$. Hence, it masks the activations of invalid actions (line 11) by computing the set of valid actions $A^{-}$ based on the current environment state (line 10). Due to the large number of possible terminal states, effective navigation of the solution space is crucial. While exploring new actions leads to better solutions, it may affect convergence and learning stability. In contrast, exploiting current knowledge (embedded in the policy network $\theta$) accelerates convergence but may lead to suboptimal results. Thus, ReLATE performs heavy exploration at the beginning of the learning process by selecting a random action with a high probability $\epsilon$ and then reduces this probability along with the learning rate $\alpha$ over time to avoid catastrophic divergence (line 18).

Once a valid action $a_{t}$ is selected, ReLATE analyzes this action to determine if an interaction with the environment is needed for reward evaluation (line 12). Specifically, ReLATE assigns a default (zero) reward for all intermediate actions $a_{t}\forall t\neq\ell(p)-1$ and only considers terminal actions $a_{\ell(p)-1}$ for reward evaluation. Most importantly, during the exploration stage, ReLATE learns to judge its own actions by gradually forming a simple reward model of the environment (omitted from Algorithm 1 for brevity), based on the actual reward of terminal actions, using a fully connected network architecture. After this predictive model reaches sufficient accuracy, the agent transitions to learning from real actions (reward evaluation) and imagined actions (reward model). When an imagined action is rated as high-value by the reward model, i.e., the estimated reward is no worse than the highest observed reward by an amount equal to the model’s error margin, the agent evaluates the real reward for that action, and updates the model accordingly. Note that even in complex environments, a simple reward predictor is sufficient, as shown in Figure 8, because it only predicts the reward of terminal actions/states, not the full transition dynamics of the target environment. Additionally, by inspecting the reward cache before computing any terminal reward, ReLATE avoids evaluating already seen action sequences.

After executing an action $a_{t}$, ReLATE stores the resulting transition in the temporary episodic memory and trains the policy network ($\theta$) via stochastic gradient descent (SGD) using a sampled minibatch from the replay buffer. During training, the value of an action selected by the policy network $\theta$ is evaluated with the target network $\theta^{-}$, which is updated at a slower rate as shown in line 19, to provide more reliable estimates of future values (Hasselt et al., 2016). At the end of every episode, ReLATE updates the best sparse encoding discovered so far based on the obtained reward from the newly constructed encoding. Nevertheless, since the reward is only generated for terminal actions/states, the agent can suffer from unstable and delayed learning as the feedback signal is both sparse and ambiguous, making it hard to determine which earlier actions mattered. To tackle this issue, ReLATE introduces a reward shaping mechanism (line 23) that allocates credit $r_{t}$ from the sparse reward $r_{\ell(p)-1}$ to all actions leading up to that reward. Despite testing several reward-shaping schemes, uniform credit distribution yielded the best results with minimal complexity, after converting speedup/slowdown factors into positive/negative rewards via a log-scale mapping. Finally, ReLATE moves the updated transitions from the temporary episodic memory to the priority-based replay buffer, which is used during training in future episodes.

4. Evaluation

We evaluate the proposed ReLATE framework for learning optimized sparse tensor encoding against the state-of-the-art, expert-designed formats, covering both mode-agnostic and mode-specific tensor representations. We characterize the performance of the target TD operations on the fifth generation of Intel Xeon Scalable processors, codenamed Emerald Rapids (EMR).

4.1. Experimental Setup

4.1.1. Test Platform

The evaluation was carried out on dual-socket Intel Xeon Platinum 8592+ processors, with EMR micro-architecture. Each EMR socket comprises 64 physical cores, running at a fixed 1.9 GHz base frequency, and has 320 MiB L3 cache. The experiments use all hardware threads ($128$) on the target processor with turbo mode disabled to ensure accurate performance measurements. The test system is equipped with 1024 GiB of DDR5 main memory, and runs Rocky Linux 9.3 distribution. The code is compiled with Intel C/C++ compiler (v2023.2.0) using the optimization flags -O3 -qopt-zmm-usage=high -xHost to fully utilize vector units. The ReLATE agent was implemented using Python 3.11 and PyTorch 2.6.0. For interactions between ReLATE and the TD environment, we use the mpi4py package (v4.0.3) built on top of the Intel MPI library (v2021.15).

4.1.2. Datasets

The experiments use all real-world sparse tensors available in the FROSTT (Smith et al., 2017) repository, with working sets that exceed cache capacity. For smaller tensors, in terms of dimensions and nonzero counts, the choice of linearized sparse encoding has negligible effect on performance. Specifically, at this small scale, ReLATE attains the same performance level as the state-of-the-art ALTO format (Laukemann et al., 2025). Table 1 lists the target sparse tensors, sorted by their number of nonzero elements (#NNZs).

4.1.3. Configurations

We compare ReLATE to the best mode-agnostic as well as mode-specific sparse tensor representations, namely linearized (Helal et al., 2021; Laukemann et al., 2025) and CSF-based (Smith et al., 2015; Smith and Karypis, 2017) formats, respectively. We use the parallel MTTKRP implementations provided by the state-of-the-art TD libraries: ALTO111Available at: https://github.com/IntelLabs/ALTO for mode-agnostic formats and SPLATT222Available at: https://github.com/ShadenSmith/splatt for mode-specific formats, which demonstrated superior performance to prior TD methods (Laukemann et al., 2025). Additionally, we leverage the ALTO library as the external TD environment in our ReLATE framework, and use the hyperparameters detailed in Table 2. During the offline training, the ReLATE agent and the external TD environment are deployed on different compute nodes. Once a policy is learned, the optimized sparse encoding discovered by ReLATE is evaluated against the expert-designed formats. All offline training tasks are subject to a $6$-hour timeout. To achieve the best performance, SPLATT is configured to use $N$ tensor copies for an $N$-dimensional tensor, and we report the best results obtained across two format variants: standard CSF and CSF with tensor tiling. Likewise, ALTO is configured to set the size of its linearized index to the smallest multiple of the native word size (i.e., $64$ or $128$ bits) appropriate for the target sparse tensors. Moreover, in the ALTO library, we set the minimum fiber reuse required for reduction-based synchronization to $8$ to enhance its performance on the target EMR processors. As in prior studies (Smith et al., 2015; Helal et al., 2021; Laukemann et al., 2025), the experiments decompose each tensor data set into $16$ rank-one components, employing double-precision arithmetic and 64-bit native integers. To benchmark TD performance, we measure the execution time required to perform parallel MTTKRP operations over all tensor modes. For tensors with fewer than a billion nonzero elements, we repeat the MTTKRP operation $10$ times to mitigate timing measurement noise.

4.2. Performance Results and Analysis

Figures 5 and 6 compare the performance and storage of the learned sparse tensor representations discovered by ReLATE to prior state-of-the-art formats (Laukemann et al., 2025). Since expert-designed formats can be highly sensitive to the data distributions of sparse tensors (Helal et al., 2021), we analyze the performance under random permutations of tensor indices (“Shuffled”) to break down the inherent data skewness, in addition to the original (non-permuted) tensors. The results demonstrate that ReLATE outperforms the best mode-agnostic as well as mode-specific formats across all sparse tensors, while using the same memory storage as the mode-agnostic ALTO format.

Specifically, ReLATE realizes a geometric-mean speedup of $1.4-1.46\times$, and up to $2\times$ speedup, over the best expert-designed format. Compared to the mode-specific SPLATT, ReLATE achieves up to $16.9\times$ speedup, with a geometric-mean speedup of $3.28-3.69\times$. The results show that ReLATE delivers substantial performance gains for both the original and randomly permuted sparse tensors, due to the adaptive nature of its learned sparse encoding. SPLATT has better average performance on the randomly permuted tensors compared to the original ones, as random permutation reduces data skewness, and in turn improves the workload balance of CSF-based formats; however, both ALTO and ReLATE significantly outperform SPLATT, thanks to the workload balance, high data reuse, and low synchronization cost of linearized tensor formats (Helal et al., 2021; Laukemann et al., 2025).

Additionally, Figure 5 and Table 1 show that the performance gains of our learned sparse encoding compared to the best expert format increase with the tensor size and sparsity. This performance trend not only highlights the limitations of prior input-agnostic heuristics for constructing sparse tensor formats, but also indicates that learning-augmented algorithms are crucial for tackling low-density, large-scale tensors, which are particularly challenging to optimize on modern parallel processors, due to their limited data reuse, low arithmetic intensity, and random memory access (Choi et al., 2018; Nguyen et al., 2022; Laukemann et al., 2025). Although the patents tensor is large, in terms of the number of nonzero elements, it has short dimensions; therefore, the performance gap between ReLATE and ALTO is small for this tensor, as its factorization data has comparable size to L$3$ cache. While deli-3d and flicker-3d have higher density than deli-4d and flicker-4d, ReLATE realizes even better performance than the expert formats in the former cases compared to the latter. Further analysis reveals that the density increase stems from removing the $4^{th}$ dimension in these tensors, which represents dates and is extremely short. Consequently, because TD operations along this mode exhibit high fiber reuse, removing it exacerbates the challenges of executing TD operations along the remaining longer modes, widening the performance gap between ReLATE and the prior input-agnostic formats.

Because the mode-agnostic ALTO library uses a single tensor copy to perform TD operations across all modes, its memory storage is significantly lower than the mode-specific SPLATT library, which maintains one tensor copy per mode to maximize performance. Since the storage of linearized tensor formats depends on the tensor shape rather than it data distribution, in contrast to SPLATT and other compressed tensor formats (Helal et al., 2021; Soh et al., 2023; Laukemann et al., 2025), ReLATE matches the storage requirements of ALTO for the original and randomly permuted tensor data sets. Hence, Figure 6 shows that SPLATT consumes more storage by up to a factor of $4.2\times$ relative to both ReLATE and ALTO. Although random permutation of sparse tensors improves workload balance and performance for SPLATT, the results indicate that it requires more memory on average to represent the randomly permuted tensors compared to the original tensors.

4.2.1. Roofline Analysis

To gain insight into the performance of ReLATE compared to the prior state-of-the-art format, we constructed a Roofline model (Williams et al., 2009) for the target EMR processors and collected performance counters using the LIKWID tool suite v5.4.1 (Gruber et al., 2024). Under the Roofline model, performance is bounded by $\Phi=\mathrm{min}(\Phi_{\mathrm{peak}},BW_{\mathrm{peak}}\times OI)$, with $\Phi_{\mathrm{peak}}$ denoting the theoretical peak compute throughput, $BW_{\mathrm{peak}}$ peak memory bandwidth, and $OI$ operational intensity (FLOPs per byte). We estimate the peak performance $\Phi_{\mathrm{peak}}$ assuming a steady state, in which a core can issue two fused multiply-add (FMA) per cycle on 512-bit vector registers, delivering up to 32 double-precision FLOPs per cycle per core. We augment our Roofline model to account for the cache bandwidth, as in prior work (Helal et al., 2021; Laukemann et al., 2025). We measure L2/L3 and main-memory bandwidth with likwid-bench from the LIKWID tool suite, and adopt the theoretical L1 bandwidth: two cache lines (512-bit transactions) per cycle per core.

Figure 7 and table 3 show the detailed performance analysis for three representative tensors. Because deli-3d is highly sparse with limited data reuse, its performance using ALTO is largely bounded by main memory bandwidth. Nevertheless, our ReLATE agent discovers a sparse encoding that significantly decreases L$3$ miss ratio, reducing main memory volume by $43\%$. Although reddit is a low-density tensor, it has a large number of nonzero elements relative to its dimensions, increasing data reuse compared to deli-3d. Consequently, the performance of reddit using ALTO is bounded by L$3$ cache bandwidth. Even under this relatively high data reuse, ReLATE improves the L$3$ hit ratio and reduces memory traffic by $41\%$. In contrast to the previous cases, patents has a higher density due to its small dimensions. Hence, its factorization data is largely served from L$2$ cache, as evidenced by the markedly lower memory volume compared to reddit, even though both tensors have a comparable number of nonzero elements. Thus, while ReLATE slightly reduces the L$2$ miss ratio, the performance impact is limited compared with other large-scale, low-density tensors.

4.3. Agent Effectiveness

To evaluate the effectiveness of our ReLATE agent, we validated its reward model and categorized its actions accordingly. Figure 8 shows the average estimation error of the reward model across the first $1000$ training episodes. The estimation error is measured as the absolute difference between the estimated and actual reward, normalized by the magnitude of the actual reward. This error is reported as an average over 10 test samples that the model had not yet encountered during training. The results indicate that the reward model can quickly learn the environment dynamics, as evidenced by the sharp decline in estimation error, settling below the error threshold after $500$ episodes at the latest.

Figure 9 illustrates the type and number of actions taken by the agent in every sparse tensor environment. The actions are categorized into real, imagined, and cached. Reward evaluation is only required for real actions. Due to the adaptive nature of our environment representation and model architecture, as detailed in $\S$3, the overall number of actions is proportional to the environment complexity, in terms of number of dimensions and dimension lengths. The results show that real actions represent a fraction ($34\%$ on average) of all actions taken by the agent because of the effectiveness of the proposed action filtering mechanisms. Model-based action filtering, in which the agent uses imagined/estimated reward, is crucial in complex environments, such as deli-3/4d, flicker-3/4d, nell-1, amazon, and reddit, while cache-based filtering is useful in simpler environments, like darpa.

The detailed analysis indicates that ReLATE was able to learn effective policies and discover efficient sparse representations that outperform the expert formats within $5000$ episodes, while operating in isolation from the target TD environment. More importantly, the worst-case performance encountered by the agent, before learning the environment dynamics, is on par with the mode-specific SPLATT and amounts to $0.2-0.5\times$ of the best format performance.

5. Conclusion

This work introduces a novel approach for optimizing sparse tensor decomposition through reinforcement learning techniques, eliminating the need for labeled training samples. By effectively learning high-quality sparse tensor representations, our ReLATE framework delivers substantial performance gains across diverse data sets, achieving up to $2\times$ speedup over traditional expert-designed formats. Our analysis reveals that the learned sparse representations address the key performance bottlenecks, yielding better cache utilization, lower memory traffic, and higher floating-point throughput, especially for large-scale, low-density tensors that are hard to optimize on modern parallel processors. Future work will investigate transferring these learning capabilities to related sparse tensor operations.