Step-TP: A Grounded, Step-Level Dataset with Chain-of-Thought Reasoning for LLM-Guided Tensor Program Optimization

• URL: http://arxiv.org/abs/2605.25954v1
• Date: Mon, 25 May 2026 15:29:49 GMT
• Title: Step-TP: A Grounded, Step-Level Dataset with Chain-of-Thought Reasoning for LLM-Guided Tensor Program Optimization
• Abstract summary: Step-TP is a post-training dataset for tensor program optimization.<n>It provides grounded, atomic, step-level supervision with structured chain-of-thought (CoT) reasoning.<n>Step-TP forms a closed reasoning loop over intermediate program states, enabling reliable multi-step optimization.
• Score: 24.150613011881774
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Despite the strong reasoning capabilities of large language models (LLMs), optimizing the execution efficiency of tensor programs remains challenging due to the need for precise, composable transformation decisions. Recent LLM-guided approaches frame tensor program optimization as an iterative decision process, but existing datasets provide only end-to-end optimized program pairs using token-inefficient representations, lacking verifiable step-level supervision and interpretability. As a result, LLMs struggle to make reliable single-step decisions in large combinatorial optimization spaces. We introduce Step-TP, a post-training dataset for tensor program optimization that provides grounded, atomic, step-level supervision with structured chain-of-thought (CoT) reasoning. Step-TP forms a closed reasoning loop over intermediate program states, enabling reliable multi-step optimization rather than outcome imitation. Its design is guided by four principles: (i) a token-efficient, verifiable intermediate representation (IR) that deterministically lowers to TVM TIR; (ii) atomic and composable optimization strategies that decompose complex trajectories into interpretable single-step decisions; (iii) structured CoT supervision coupled with explicit IR-to-IR state transitions; and (iv) strategy filtering to balance coverage while preventing shortcut exploitation. The dataset and implementation are available at a GitHub link, https://github.com/LIUMENGFAN-gif/StepTP.

Related papers

• HiDiffTIR: Hierarchical Difficulty-Aware Policy Optimization for Multi-Turn Tool-Integrated Reasoning [95.29291773706224]
  Tool-Integrated Reasoning is a capability for LLM agents to solve complex tasks by interacting with external tools iteratively.<n>Existing approaches typically assign uniform trajectory-level advantages and treat all correct tool calls equally.<n>This can lead to imprecise learning signals that do not adequately distinguish between trivial and challenging tool-use patterns.
  arXiv  Detail & Related papers  (2026-08-22T09:14:15Z)
• ThinkBooster: A Unified Framework for Seamless Test-Time Scaling of LLM Reasoning [110.46431027868954]
  ThinkBooster is a framework for seamless test-time compute scaling of large language model (LLM) reasoning.<n>It consists of a modular Python library implementing state-of-the-art TTC scaling strategy and scorer families.<n>A deployable OpenAI-compatible proxy service enables drop-in integration of adaptive reasoning into real-world applications.
  arXiv  Detail & Related papers  (2026-06-05T05:28:46Z)
• Execution-Verified Reinforcement Learning for Optimization Modeling [49.171122807323634]
  Execution-verified learning framework treats a mathematical programming solver as a deterministic, interactive verifier.<n>Experiments on NL4OPT, MAMO, IndustryOR, and OptiBench across Gurobi, OR-Tools, and COPT show that EVOM matches or outperforms process-supervised SFT.
  arXiv  Detail & Related papers  (2026-04-01T03:39:11Z)
• From Brute Force to Semantic Insight: Performance-Guided Data Transformation Design with LLMs [48.83701310501069]
  Large language models (LLMs) have achieved notable performance in code synthesis.<n>We introduce a performance-aware, closed-loop solution that enables LLMs to autonomously engineer optimal transformations.<n>We fine-tune LLMs with Low-Rank Adaptation on a novel repository of more than 6,000 empirically evaluated PyTorch augmentation functions.
  arXiv  Detail & Related papers  (2026-01-07T11:13:02Z)
• Adaptive Dependency-aware Prompt Optimization Framework for Multi-Step LLM Pipeline [9.013236765328301]
  We propose ADOPT, an Adaptive Dependency-aware Prompt Optimization framework for multi-step LLM pipelines.<n> ADOPT explicitly models the dependency between each LLM step and the final task outcome, enabling precise text-gradient estimation.<n>Experiments on real-world datasets and diverse pipeline structures show that ADOPT is effective and robust.
  arXiv  Detail & Related papers  (2025-12-31T15:46:37Z)
• Rectifying LLM Thought from Lens of Optimization [48.98086817378953]
  Long chain-of-thought (CoT) prompting enables thorough exploration and deliberation.<n>Despite advances, long-CoT LLMs often exhibit suboptimal reasoning behaviors.<n>We introduce RePro, a novel approach to refine LLM reasoning during post-training.
  arXiv  Detail & Related papers  (2025-12-01T17:41:08Z)
• PROPA: Toward Process-level Optimization in Visual Reasoning via Reinforcement Learning [30.44007644340425]
  We introduce PROPA, a novel framework that integrates Monte Carlo Tree Search (MCTS) with GRPO to generate dense, process-level rewards and optimize reasoning at each intermediate step without human annotations.<n>Across seven benchmarks and four VLM backbones, PROPA consistently outperforms both SFT- and RLVR-based baselines.<n>It achieves up to 17.0% gains on in-domain tasks and 21.0% gains on out-of-domain tasks compared to existing state-of-the-art.
  arXiv  Detail & Related papers  (2025-11-13T13:06:12Z)
• CoT Vectors: Transferring and Probing the Reasoning Mechanisms of LLMs [33.63911145333626]
  Chain-of-Thought prompting has emerged as a powerful approach to enhancing the reasoning capabilities of Large Language Models.<n>Existing implementations, such as in-context learning and fine-tuning, remain costly and inefficient.<n>We introduce CoT Vectors, compact representations that encode task-general, multi-step reasoning knowledge.
  arXiv  Detail & Related papers  (2025-10-01T06:58:23Z)
• SymRTLO: Enhancing RTL Code Optimization with LLMs and Neuron-Inspired Symbolic Reasoning [30.938876549335067]
  This paper presents SymRTLO, a novel neuron-symbolic RTL optimization framework.<n>A symbolic module is proposed for analyzing and optimizing finite state machine (FSM) logic.<n> Experiments on the RTL-Rewriter benchmark with Synopsys Design Compiler and Yosys show that SymRTLO improves power, performance, and area (PPA) by up to 43.9%, 62.5%, and 51.1%, respectively.
  arXiv  Detail & Related papers  (2025-04-14T16:15:55Z)
• In-context Demonstration Matters: On Prompt Optimization for Pseudo-Supervision Refinement [71.60563181678323]
  Large language models (LLMs) have achieved great success across diverse tasks, and fine-tuning is sometimes needed to further enhance generation quality.<n>To handle these challenges, a direct solution is to generate high-confidence'' data from unsupervised downstream tasks.<n>We propose a novel approach, pseudo-supervised demonstrations aligned prompt optimization (PAPO) algorithm, which jointly refines both the prompt and the overall pseudo-supervision.
  arXiv  Detail & Related papers  (2024-10-04T03:39:28Z)
• Towards Explainable Evolution Strategies with Large Language Models [0.0]
  This paper introduces an approach that integrates self-adaptive Evolution Strategies (ES) with Large Language Models (LLMs) By employing a self-adaptive ES equipped with a restart mechanism, we effectively navigate the challenging landscapes of benchmark functions. An LLM is then utilized to process these logs, generating concise, user-friendly summaries.
  arXiv  Detail & Related papers  (2024-07-11T09:28:27Z)
• LLM as a Complementary Optimizer to Gradient Descent: A Case Study in Prompt Tuning [69.95292905263393]
  We show that gradient-based and high-level LLMs can effectively collaborate a combined optimization framework.<n>In this paper, we show that these complementary to each other and can effectively collaborate a combined optimization framework.
  arXiv  Detail & Related papers  (2024-05-30T06:24:14Z)
• Intuitive Fine-Tuning: Towards Simplifying Alignment into a Single Process [19.986235452236272]
  Supervised Fine-Tuning (SFT) and Preference Optimization (PO) are key processes for aligning Language Models (LMs) with human preferences post pre-training.<n>We introduce Intuitive Fine-Tuning (IFT) to integrate SFT and PO into a single process.<n>IFT performs comparably or even superiorly to SFT and some typical PO methods across several tasks.
  arXiv  Detail & Related papers  (2024-05-20T08:23:28Z)
• Goal Kernel Planning: Linearly-Solvable Non-Markovian Policies for Logical Tasks with Goal-Conditioned Options [54.40780660868349]
  We introduce a compositional framework called Linearly-Solvable Goal Kernel Dynamic Programming (LS-GKDP)<n>LS-GKDP combines the Linearly-Solvable Markov Decision Process (LMDP) formalism with the Options Framework of Reinforcement Learning.<n>We show how an LMDP with a goal kernel enables the efficient optimization of meta-policies in a lower-dimensional subspace defined by the task grounding.
  arXiv  Detail & Related papers  (2020-07-06T05:13:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.