Research map / Architectures and efficiency / Pruning and efficient tuning
COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes COGNATE, a transfer-learning cost model for sparse tensor programs that pre-trains on cheap CPU data and few-shot fine-tunes on expensive emerging accelerators (SPADE simulator, A100 GPU).
Problem
Learned cost models for optimizing sparse tensor programs (SpMM, SDDMM) on early-stage accelerators are hard to train because simulator-based data collection is extremely expensive, and naive CPU-to-accelerator transfer fails due to heterogeneous configuration spaces.
Approach
Modify WACO's cost model with a configuration mapper that approximately maps comparable CPU/SPADE/GPU optimizations (strip-mining, reordering) into a shared space, per-platform autoencoders that compress hardware-specific parameters into fixed-size latent codes, and a shallower multi-scale input featurizer; pre-train on CPU and few-shot fine-tune on the target.
Key results
- On SPADE, COGNATE (Top-5) reaches 1.47x average SpMM speedup (up to 5.46x) and 1.39x for SDDMM, outperforming WACO+FM (1.09x) by 28.44%.
- On an A100 GPU, COGNATE (Top-5) achieves 1.17x average SpMM speedup versus 1.01x for cuSPARSE.
- Fine-tuning with only 5 SPADE matrices gives 1.40x speedup, comparable to 1.42x with 1000 matrices of target data, using about 5% of the data needed by accelerator-only models.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use COGNATE Top-5 configuration selection instead of Top-1 | COGNATE on SPADE SpMM · SPADE SpMM | speedup: 1.40× -> 1.47× | Helped |
| Use COGNATE instead of WACO+FM | SPADE SpMM · SPADE SpMM | speedup: 1.09× -> 1.40× | Helped |
| Add the input featurizer (ablation) | COGNATE · SPADE SpMM | speedup: 1.26x -> 1.40x | Helped |
| Add the configuration mapper (ablation) | COGNATE · SPADE SpMM | speedup: 1.16x -> 1.40x | Helped |
| Add the latent encoder (ablation) | COGNATE · SPADE SpMM | speedup: 1.01x -> 1.40x | Helped |
| Replace the MLP predictor with a Transformer predictor (ablation) | COGNATE · SPADE SpMM | speedup: 1.40× -> 1.36× | Hurt |
| Fine-tune using 3 rather than 5 matrices (ablation) | COGNATE · SPADE SpMM | speedup: 1.40× -> 1.30× | Hurt |
| Increase fine-tuning data from 5 to 7 matrices (ablation) | COGNATE · SPADE SpMM | speedup: 1.40× -> 1.41× | No clear effect |
| Increase fine-tuning data from 5 to 1,000 matrices (ablation) | COGNATE · SPADE SpMM | speedup: 1.40× -> 1.42× | No clear effect |
| Increase source training data from 100 to 1,000 matrices (ablation) | COGNATE with SPADE fine-tuning on 5 matrices · SPADE SpMM | speedup: 1.40 -> 1.19 | Hurt |
Surprising findings
- Training the source model on more CPU data (1000 matrices) yields worse transfer, attributed to overfitting to source-specific characteristics (negative transfer).
- Zero-shot transfer from the CPU model falls below baseline on both SPADE (0.71x) and GPU.
Limitations
- Evaluated only on two target platforms (SPADE simulator and A100 GPU); PIUMA and Vesper could not be tested.
- Only SpMM and SDDMM are evaluated; search space is constrained to 256 configurations for SPADE.
- Some matrices show the baseline outperforming COGNATE; GPU zero-shot transfer performs poorly.
Experimental setup
Datasets: SuiteSparse Matrix Collection
Benchmarks: SuiteSparse Matrix Collection (715 test matrices)
Metrics: geomean speedup top-1 and top-5 speedup APE OPA Kendall's Tau data collection overhead
Baselines: WACO+FA WACO+FM cuSPARSE zero-shot CPU model no-transfer model
Builds on: WACO TACO SparseTIR MinkowskiEngine
Closest papers by meaning
- Harnessing Manycore Processors with Distributed Memory for Accelerated Training of Sparse and Recurrent Models (AAAI 2024 · similarity 0.84)
- SpeedLoader: An I/O efficient scheme for heterogeneous and distributed LLM operation (NeurIPS 2024 · similarity 0.84)
- Sparse Cocktail: Every Sparse Pattern Every Sparse Ratio All At Once (ICML 2024 · similarity 0.84)
- Tensorion: A Tensor-Aware Generalization of the Muon Optimizer (NeurIPS 2026 · similarity 0.84)
- TSENOR: Highly-Efficient Algorithm for Finding Transposable N:M Sparse Masks (NeurIPS 2025 · similarity 0.84)
- Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage (NeurIPS 2025 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.