Research map / Architectures and efficiency / RNNs and normalization
HiRemate: Hierarchical Approach for Efficient Re-materialization of Neural Networks
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 HiRemate, a hierarchical re-materialization framework that recursively partitions a general data-flow graph and combines a new hierarchical ILP (H-ILP) with existing solvers to produce memory-budgeted PyTorch training schedules.
Problem
Re-materialization (recomputation) for training DNNs under GPU memory budgets is NP-hard on general graphs; existing exact ILP methods (Checkmate) do not scale beyond ~70-90 nodes, and chain-based methods (Rockmate, Rotor) lack generality for long skip connections.
Approach
Recursively partition the computation graph into convex subgraphs (H-Partition), solve bottom-level subgraphs under multiple budgets with different solvers, and recombine their options at higher levels with a new ILP (H-ILP) using phantom nodes and option-specific memory corrections, yielding a global schedule.
Key results
- Achieves up to 50-70% memory reduction with only 10-15% iteration-time overhead, per the abstract.
- On a 4-layer encoder-decoder Transformer, Rockmate takes 15 hours to solve while H-ILP achieves better results within 12 minutes.
- Table 2: H-ILP overhead closely matches RK-Checkmate optimal (e.g., UNet 5.44% vs 5.32% at 7.1 GB; 5-layer GPT2 9.24% vs 6.02% at 1.25 GB).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use H-ILP instead of RK-CHECKMATE | HIREMATE · UNet, 7.1 GB budget | iteration time overhead: 5.32% -> 5.44% | Hurt |
| use H-ILP instead of RK-CHECKMATE | HIREMATE · 2-layer Transformer, 7.0 GB budget | iteration time overhead: 5.32% -> 5.34% | Hurt |
| use H-ILP instead of RK-CHECKMATE | HIREMATE · MLP-Mixer, 6.0 GB budget | iteration time overhead: 2.47% -> 2.44% | Helped |
| use H-ILP instead of RK-CHECKMATE | HIREMATE · 5-layer GPT2, 1.25 GB budget | iteration time overhead: 6.02% -> 9.24% | Hurt |
Surprising findings
- Increasing hierarchy depth causes only slight solution-quality degradation for H-ILP, despite harder subproblems.
Limitations
- Not adapted to dynamic neural network architectures with input-dependent control flow.
- Current H-partition implementation has O(N^2 log N) complexity, requiring more work for very large graphs.
- Does not integrate offloading or pipeline parallelism.
Experimental setup
Metrics: peak memory iteration time solving time relative overhead
Baselines: Checkmate TW-Remat Moccasin Rockmate PyTorch autodiff
Builds on: Checkmate RK-Checkmate Rockmate Rotor TW-Remat
Closest papers by meaning
- Efficient Rematerialization for Deep Networks (NeurIPS 2019 · similarity 0.87)
- Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re-materialization in PyTorch (ICML 2023 · similarity 0.87)
- DOGE-Train: Discrete Optimization on GPU with End-to-End Training (AAAI 2024 · similarity 0.83)
- Auto-reconfiguration for Latency Minimization in CPU-based DNN Serving (ICML 2025 · similarity 0.83)
- StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning (ICML 2026 · similarity 0.83)
- SpeedLoader: An I/O efficient scheme for heterogeneous and distributed LLM operation (NeurIPS 2024 · 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.