RED-PIM: Reducing Data Movement for Transformers using Processing-in-Memory

Source: arXiv:2607.21731 · Published 2026-07-23 · By Zahra Yousefijamarani, Alaa Alameldeen

TL;DR

This work tackles the fundamental performance bottleneck of transformer inference caused by extensive data movement between memory and processing units during self-attention operations. Standard transformer attention mechanisms require computing and storing large N×N attention matrices, which exceed typical memory bank capacities and induce costly inter-bank communication. Prior Processing-In-Memory (PIM) solutions mitigate some data movement but still suffer from quadratic communication overhead and scalability issues for long sequences.

The authors propose RED-PIM, an algorithm-architecture co-design that reduces inter-bank data movement from O(N^2) to O(N) by adopting an alternative attention formulation that replaces the full attention matrix with compact d×d intermediate matrices (where d ≪ N). RED-PIM also reorganizes computation to keep operations local within memory banks, introduces hierarchical aggregation for communication efficiency, and devises a hardware-friendly softmax approximation to accelerate key operations. Evaluations on real-world datasets and transformer workloads demonstrate geometric mean inference time reductions of 66.42% over baseline PIM methods, with up to 99.6% improvement on long documents, while maintaining or improving accuracy. These results show RED-PIM’s practical effectiveness in enabling scalable, low-latency transformer inference on PIM architectures.

Key findings

• RED-PIM reduces inter-bank data movement complexity from O(N^2) in baseline PIM implementations to O(N), enabling more scalable attention computations.
• Inference time reductions range from 16.05% to 99.99% across evaluated workloads, with a geometric mean improvement of 66.42% over baseline PIM implementations.
• On long-document datasets (IMDB, PubMed, Arxiv, WikiHop, GovReport), RED-PIM achieves a 99.60% geometric mean performance improvement.
• On shorter-document datasets (GLUE benchmark), RED-PIM yields a 13.44% average speedup while maintaining or improving accuracy.
• RED-PIM shrinks intermediate attention matrix size from N×N to d×d, significantly reducing memory footprint per attention layer.
• The hardware-friendly softmax approximation using bitwise operations and a small lookup table achieves comparable accuracy (relative error ~2.5×10^-3) while reducing LUT size by more than 600× compared to naive LUT methods.
• Hierarchical transfer strategy for inter-bank data aggregation reduces communication overhead and latency compared to naive centralized approaches.
• RED-PIM’s data movement analysis and architectural optimizations enable efficient support of sequences longer than typical bank memory capacities.

Methodology — deep read

1. Threat model & assumptions: The adversary is not explicitly defined in this architecturally focused work, but the goal is to optimize transformer inference efficiency under the assumption that large sequences (N up to thousands) and large model dimensions (d) need to be processed on Processing-In-Memory (PIM) hardware with limited memory bank capacity (typical HBM bank size ~32MB). The baseline assumes weight matrices (WQ, WK, WV) and tokens are distributed across memory banks.

2. Data: They evaluate transformer inference workloads using real-world natural language datasets including GLUE benchmark tasks for shorter sequences and longer document datasets like IMDB, PubMed, Arxiv, WikiHop, and GovReport. The paper lists datasets used for evaluation but does not specify exact training splits since this is an inference-focused study.

3. Algorithm / Architecture:

• Baseline PIM implementation: standard Scaled Dot-Product Attention (SDPA) where Query-Key multiplication produces an N×N attention matrix S distributed across banks, resulting in O(N^2/m) inter-bank communication.
• RED-PIM algorithm: Replaces explicit computation of full attention matrix with alternative decomposition (adapted from Zhuoran et al. [36]) where softmax is applied separately to Q and K^T, then multiplied with Value matrices. This produces an intermediate d×d matrix L per bank.
• Architectural design: Utilizes FIMDRAM PIM architecture with local compute units (PCUs) in each memory bank, supporting FP16 vector multiplications and specialized instructions.
• Hierarchical data transfer mechanism aggregates partial computations locally within bank groups, followed by channel-level and die-level aggregation to reduce inter-bank traffic and latency.
• Hardware-friendly softmax: Implements exp(x) via identity e^x = 2^α exploiting fast bit-shift for integer part and a small LUT with linear interpolation for fractional part, trading off minimal accuracy for significant memory and compute savings.

4. Training regime: Since the paper focuses on inference acceleration, model training details are inherited from pretrained transformers on standard datasets. The paper does not discuss retraining but verifies accuracy preservation on evaluation benchmarks.

5. Evaluation protocol:

• Metrics: Inference time (latency) reduction compared to baseline PIM transformer acceleration.
• Benchmarks: Multiple datasets covering short and long sequences.
• Simulations use a custom extension of the DAMOV simulator to model inter-bank data movement and PIM computation.
• Ablations include testing different data movement strategies and softmax implementations.
• Accuracy validation confirms that approximation techniques do not degrade model performance.

6. Reproducibility:

• A customized simulation framework based on DAMOV is developed but the paper does not explicitly mention public release of code or models.

Concrete example: For a sequence of length N distributed across memory banks (each holding m tokens), baseline attention requires moving Q vectors across N/m banks to multiply with all K^T columns, causing O(N^2/m) data movements. RED-PIM instead computes local softmax of Q and K^T separately, producing compact d×d matrices locally, which are then hierarchically aggregated across banks, reducing data movement complexity to O(N/m). Using the optimized softmax and hierarchical transfer, RED-PIM achieves large inference speedups on long sequences while fitting data within single banks, avoiding costly inter-bank transfers.

Technical innovations

• Algorithmic re-formulation of self-attention to avoid explicit full N×N attention matrix computation by decomposing softmax operations and multiplying smaller d×d intermediate matrices.
• Hierarchical inter-bank data transfer and aggregation mechanism that reduces communication overhead and latency compared to naive centralized aggregation.
• Hardware-optimized softmax implementation using bitwise exponentiation and a small lookup table that reduces memory usage by >600× while maintaining relative error ~2.5×10^-3.
• Integration of the above algorithmic improvements with the FIMDRAM PIM architecture, enabling efficient in-bank vector multiplications and scalable transformer inference on memory-constrained banks.

Datasets

• GLUE benchmark suite — various natural language understanding tasks — public
• IMDB — long document sentiment classification — public
• PubMed — biomedical scientific papers — public
• Arxiv — scientific papers — public
• WikiHop — reading comprehension dataset — public
• GovReport — document classification — public

Baselines vs proposed

• Baseline PIM transformer implementation: inference time reductions from 0% baseline to RED-PIM 16.05% to 99.99%, geometric mean 66.42%
• On long documents (IMDB, PubMed, Arxiv, WikiHop, GovReport): baseline vs RED-PIM geometric mean speedup 99.60%
• On short documents (GLUE benchmark): baseline vs RED-PIM 13.44% average speedup
• Softmax approximation methods: simple LUT relative error ~2.5×10^-3 using 312.5KB LUT vs RED-PIM method relative error ~2.5×10^-3 using 0.5KB LUT vs Taylor series poor accuracy (2.6×10^77 error)

Figures from the paper

Figures are reproduced from the source paper for academic discussion. Original copyright: the paper authors. See arXiv:2607.21731.

Fig 1: The Transformer model architecture [42].

Fig 2: d illustrates the internal architecture of a PCU.

Fig 3: Token and weight matrix mapping across banks.

Fig 4: Regular Scaled Dot-Product Attention mechanism. Here, N denotes the sequence length, d is the embedding dimension, and m represents the maximum

Fig 5: RED-PIM Scaled Dot-Product Attention. Parts highlighted in red indicate portions stored locally within the same bank. σ denotes the softmax

Fig 6: illustrates the hierarchical transfer process using

Fig 7 (page 5).

Fig 8 (page 5).

Limitations

• Evaluation is performed via simulation rather than on physical PIM hardware, which may affect real-world performance estimations due to unmodeled overheads.
• The approach assumes static and distributed placement of weight matrices and tokens; dynamic memory management and model updates are not addressed.
• Accuracy evaluation focuses on inference tasks; effects on training or fine-tuning are not analyzed.
• The algorithmic reformulation is adapted from prior efficient attention methods but may have limitations on certain transformer variants or attention heads.
• Scalability to extreme model sizes (multiple billions of parameters) requires further validation, especially regarding weight tiling and data movement overheads.

Open questions / follow-ons

• How does RED-PIM perform on other transformer variants such as sparse attention or memory-augmented models?
• Can the RED-PIM framework be extended effectively for transformer training workloads, beyond inference?
• What are the trade-offs between latency, energy consumption, and accuracy when deploying RED-PIM on real silicon prototypes?
• How does RED-PIM integrate with optimized transformer architectures involving pruning, quantization, or model compression?

Why it matters for bot defense

Bot-defense and CAPTCHA inference systems increasingly rely on transformer-based models for user behavior analysis, anomaly detection, and natural language understanding. RED-PIM’s algorithm-architecture co-design offers a compelling approach to efficiently execute large-scale transformer inference directly within memory, significantly reducing latency and energy associated with data movement. This is particularly critical for real-time bot detection scenarios where decisions must be made promptly on long sequences of user interactions or log data. Practitioners in CAPTCHA and bot defense engineering should consider RED-PIM’s data movement reduction strategies and hierarchical aggregation mechanisms when designing low-latency transformer accelerators embedded in edge or server hardware. The hardware-friendly softmax approximation may also enable more efficient deployment on specialized memory devices, enhancing throughput without sacrificing accuracy. Overall, RED-PIM demonstrates an avenue for scalable transformer inference that may help maintain responsiveness and accuracy in security-sensitive ML-powered bot-defense systems.

Cite

@article{arxiv2607_21731,
  title={ RED-PIM: Reducing Data Movement for Transformers using Processing-in-Memory },
  author={ Zahra Yousefijamarani and Alaa Alameldeen },
  journal={arXiv preprint arXiv:2607.21731},
  year={ 2026 },
  url={https://arxiv.org/abs/2607.21731}
}

Read the full paper

• arXiv:2607.21731 (abstract)
• arXiv:2607.21731 (PDF)