Optimizing Tensor Train Decomposition in DNNs for RISC-V Architectures Using Design Space Exploration and Compiler Optimizations

AI summary

Overview

Research area: Deep neural network compression (low-rank factorization) combined with design space exploration and compiler-level optimization for RISC-V edge processors.

Technical level: Advanced. The paper assumes familiarity with tensor decompositions, tensor contraction (Einsum), FLOP/parameter counting, polyhedral compilation, and embedded ISA constraints.

Scope: The paper proposes a layer-level methodology and tool that prune the Tensor Train Decomposition design space of fully connected layers and then apply RISC-V-specific compiler optimizations to the surviving candidate solutions.

What This Paper Is About

Fully connected (FC) layers dominate the parameter count, FLOPs, memory use, and inference time of many DNNs, including CNNs and large language models, which makes them hard to deploy on resource-constrained RISC-V platforms. Low-rank factorization with Tensor Train Decomposition (TTD, via the TensorFlow T3F library) can shrink those layers, but the space of possible decomposition shapes and rank lists is astronomically large — the authors note it can reach up to 10^33 possible solutions for larger layers, and even a tiny 120×84 FC layer has a large design space. The goal is to navigate that space efficiently, without exhaustively evaluating every configuration, and to produce optimized code that actually runs fast on RISC-V.

Key Contributions

1. A new method for efficiently reducing the design space in Tensor Train Decomposition, enabling faster and more effective exploration of factorized solutions.
2. A compiler optimization approach proposed to accelerate custom T3F layers of TensorFlow on RISC-V processors.
3. A comprehensive experimental evaluation conducted on seven CNNs and six LLMs, showing that the method reduces the exploration space and achieves substantial speedups for custom T3F layers.
4. A holistic, layer-level DSE methodology that combines high-level design space reduction (FLOPs/memory trade-offs), low-level design space reduction (inference efficiency heuristics), and RISC-V-tailored code optimization — described as the first work to integrate LRF compression with compiler optimization for RISC-V.

Main Findings

• Speedups over existing toolchains: The custom TT-decomposed layers run on average 3× faster than IREE and 8× faster than Pluto on the same compressed model (stated in the abstract as "3× faster than IREE and 8× faster than Pluto," and in the introduction as an average inference time reduction of three times versus IREE and eight times versus Pluto).
• Design space reduction: The proposed methodology reduces the design space by several orders of magnitude.
• Shape alignment reduces FLOPs: The proposed "aligned" combination shape — where n1 ≤ n2 ≤ … ≤ nd and m1 ≥ m2 ≥ … ≥ md — always achieves the minimum FLOPs value and a memory value close to the minimum across the six configurations shown in Figures 5 and 6.
• FLOPs and execution time do not always align: For LRF solutions with comparable memory footprints, FLOPs and inference execution time can differ significantly, which the authors use to justify evaluating configurations for computational performance and not just memory efficiency.
• Design space scale: A very small FC layer of size 120×84 already produces a large number of possible LRF configurations, and larger layers can have up to 10^33 possible solutions.
• FC layers dominate cost: Figure 1 shows FC layers dominate parameters and FLOPs across many DNN models relative to non-FC parts.
• No vendor library support for T3F: The authors state that no vendor-optimized library currently supports the T3F library, since T3F's primary objective is facilitating LRF rather than optimizing execution time.
• Accuracy is not evaluated here: Accuracy analysis after factorization is explicitly out of scope and left for future work; the methodology instead returns a list of candidate solutions per FC layer so alternatives can be selected if an accuracy constraint is not met.

Methodology in Plain English

The approach works layer by layer on FC layers and takes two inputs: the shape of the target FC layer and the details of the target RISC-V hardware. It has three stages.

1. High-level design space pruning (FLOPs vs. memory). Rather than computing exact FLOPs and parameter counts for every reshaping combination and every rank value — which the authors say is not practical — the method prunes using a structural rule. Each of the two dimensions of the weight matrix is factorized into d sub-dimensions, and the order in which those sub-dimensions appear changes the FLOP count. The authors prove that a given m_s appears in s of the FLOP summands and a given n_s appears in d−s+1 of them. From this they propose aligning the input and output shapes so that output dimensions decrease (m1 ≥ m2 ≥ … ≥ md) while input dimensions increase (n1 ≤ n2 ≤ … ≤ nd), which multiplies small values of n by small values of m and minimizes total FLOPs. This pruning is done without calculating exact memory and FLOPs values.

2. Low-level design space pruning (inference efficiency). From the surviving combinations, heuristics identify and remove solutions that fail to achieve low inference times on the target RISC-V platform, leaving only configurations likely to be fast in practice.

3. Compiler optimization of the custom T3F layers. The remaining TTD solutions are optimized with compiler techniques to reduce inference time and improve computational efficiency on RISC-V.

The output for each FC layer is a list of solutions rather than a single one, which lets a model-level workflow swap in alternatives if accuracy targets are not met. The methodology is layer-level so it can be composed with other model-level optimization techniques, and the authors state it can be extended to other processor families.

Why This Matters

Impact on research: The paper connects two areas that are usually treated separately — LRF design space exploration and hardware-specific code optimization. Prior automated rank-selection work (for example VBMF and Bayesian optimization approaches) focuses on layer weight tensors while neglecting target-device constraints and code-level optimizations; prior DSE work from the same group was limited to CNNs and did not handle the vast FC layers of large models; and work on transformer models remained specific to a single model type. This paper addresses all three gaps at once, and it is described as the first to combine LRF compression with compiler optimization tailored to RISC-V.

Real-world applications:

• Deploying compressed DNNs on edge and embedded devices built on the open RISC-V instruction set architecture.
• Natural language processing and large language models, whose transformer layers are largely FC layers and therefore dominate resource consumption.
• Autonomous systems and autonomous driving, which need low-latency inference under tight memory and energy budgets.
• Speech recognition and image processing on devices with limited processing capability, memory, and energy.

Industry relevance: RISC-V has become a popular open-source ISA for embedded and edge devices, creating demand for optimized DNN implementations. Because no vendor-optimized library currently supports T3F, toolchains in this space (IREE, Pluto, GLOW, TFLM, and hand-tuned libraries such as OneDNN, cuDNN, and BLIS) leave the T3F inference path unoptimized. The reported 3× speedup over IREE and 8× over Pluto on the same compressed model represent a concrete efficiency gain for that gap, and the layer-level design means the method can be plugged into existing model-level compression pipelines.

Future Directions

• Accuracy analysis. The paper explicitly states that a detailed accuracy analysis lies beyond the current scope and is left for future work. Quantifying accuracy degradation and validating fine-tuning or retraining strategies for the proposed solutions is the most direct next step.
• Model-level integration and end-to-end evaluation. Although the methodology is designed to compose with model-level approaches and produces per-layer lists of candidate solutions, the reported evaluation is at the layer level; end-to-end behavior across whole networks remains to be characterized.
• Extension to other processor families. The authors state the methodology can be extended to other processor families beyond RISC-V, which raises the question of how much of the low-level pruning and compiler optimization is architecture-specific.
• Broader decomposition and layer coverage. The paper focuses on TTD via T3F for FC layers; whether the pruning logic and compiler strategy transfer to other LRF methods (CP decomposition, Tucker decomposition, SVD) or to non-FC layer types is not addressed.
• Broader model coverage. Existing DSE work is largely CNN-specific or single-model-type; this work evaluates seven CNNs and six LLMs, leaving open how the approach scales to other architectures and to even larger parameter spaces than the 10^33 figure cited.

Target Audience

Researchers and engineers working on model compression, low-rank factorization, and design space exploration for deep neural networks; embedded and edge systems developers targeting RISC-V; compiler and code-generation researchers interested in tensor contraction kernels (Einsum) and polyhedral-style loop optimization; and practitioners who need to deploy DNNs, particularly transformer-based models, on devices with strict memory, latency, and energy constraints. Readers should be comfortable with tensor mathematics and performance modeling to get the most from the FLOPs and parameter derivations.

Authors’ abstract

Deep neural networks (DNNs) have become indispensable in many real-life applications like natural language processing, and autonomous systems. However, deploying DNNs on resource-constrained devices, e.g., in RISC-V platforms, remains challenging due to the high computational and memory demands of fully connected (FC) layers, which dominate resource consumption. Low-rank factorization (LRF) offers an effective approach to compressing FC layers, but the vast design space of LRF solutions involves complex trade-offs among FLOPs, memory size, inference time, and accuracy, making the LRF process complex and time-consuming. This paper introduces an end-to-end LRF design space exploration methodology and a specialized design tool for optimizing FC layers on RISC-V processors. Using Tensor Train Decomposition (TTD) offered by TensorFlow T3F library, the proposed work prunes the LRF design space by excluding first, inefficient decomposition shapes and second, solutions with poor inference performance on RISC-V architectures. Compiler optimizations are then applied to enhance custom T3F layer performance, minimizing inference time and boosting computational efficiency. On average, our TT-decomposed layers run 3x faster than IREE and 8x faster than Pluto on the same compressed model. This work provides an efficient solution for deploying DNNs on edge and embedded devices powered by RISC-V architectures.