Zen-Attention: NPU Optimization & Rare Event Detection

• Zen-Attention is a dual-purpose approach that optimizes transformer attention on AMD NPUs and refines temporal focus in rare event detection using KamNet.
• The compiler framework dynamically fuses multi-step attention operations into a single operator, reducing DRAM traffic and achieving up to 4× speedup in models like ViT and CLIP.
• In KamLAND-Zen, Zen-Attention selectively weights ns-wide PMT hit maps to improve background rejection and boost half-life sensitivity by over 10%.

Zen-Attention encompasses both an advanced compiler framework for optimizing transformer attention on AMD neural processing units (NPUs) and a neural network attention mechanism for rare event detection in the KamLAND-Zen experiment. While these approaches differ in purpose and scope, both employ the term “Zen-Attention” to denote mechanisms that dramatically improve the efficiency or interpretability of attention operations within deep learning systems (Deshmukh et al., 25 Aug 2025, Li et al., 2022).

1. Definition and Scope

Zen-Attention refers to two closely related but technologically distinct constructs:

• In high-performance inference, Zen-Attention is a compiler and runtime framework for dynamically folding the transformer attention block using spatial tiling, reduction streams, and buffer management across the AMD XDNA NPU memory hierarchy. Its central goal is to minimize DRAM bandwidth utilization and latency during the execution of attention layers on edge NPUs, targeting transformer-based models commonly bottlenecked by off-chip memory.
• In the context of spatiotemporal event detection, Zen-Attention specifies a temporal attention mechanism within KamNet—a neural network for rare event search—allowing the model to selectively weight 1.5 ns–wide time slices of photomultiplier tube (PMT) hit maps for improved background rejection and enhanced interpretability (Li et al., 2022).

2. Zen-Attention for Transformer Acceleration on AMD NPUs

The Zen-Attention compiler framework addresses the “memory-bound” nature of transformer attention, where large $Q$, $K$, and $V$ tensors require heavy DRAM traffic, dominating latency and energy consumption on edge devices such as laptops or gaming consoles. Unlike conventional approaches with hardware-managed caches, the AMD XDNA NPU employs a hierarchy of software-managed scratchpad memories (L1, L2), placing the burden of data orchestration, tiling, and prefetching on the software (Deshmukh et al., 25 Aug 2025).

2.1 Dynamic Attention Folding

Zen-Attention fuses the standard sequence of attention kernels—

1. $A = Q K^\top$
2. $A = A + B + M$ (addition of bias and mask)
3. $S = \mathrm{softmax}(A)$
4. $Z = S V$

—into a single operator that streams sub-volumes of these tensors from DRAM once, caching and reusing them across the NPU’s scratchpad array.

Tiling Scheme: Tiles along the sequence and embedding dimensions are chosen such that all necessary sub-blocks fit within the combined L1 scratchpad of multiple cores. Padding is applied as necessary to satisfy kernel unrolling requirements.

Spatial Reduction: Softmax denominators or QK dot-products are reduced across columns of cores using dedicated cascade streams—e.g., an 8-way tree for an 8-core column—thereby minimizing DRAM accesses for partial result aggregation.

2.2 Design Space Exploration and Compiler Algorithm

The framework systematically explores:

• Folding levels: Degree of operator fusion (up to three sub-layers).
• Tile sizes: Constraints dictated by L1 aggregate capacity, kernel unroll factors, and DMA alignment.
• Tensor layout: Selection among row-major, column-major, and block layouts to optimize memory strides.
• Masking/padding: Strategies for handling variable-length sequences and masking requirements.
• Interconnect mapping: Assigning attention heads and tiles to NPU columns for optimal resource use.

A pseudo-code sketch of the core “Tiler” algorithm is:

$K$1 The selected configuration is then emitted as a single fused node in the compiler’s IR, with a backend codelet issuing all required loads, reductions, transposes, and softmax computations (Deshmukh et al., 25 Aug 2025).

3. Empirical Results and Impact

Zen-Attention delivers substantial gains in transformer inference latency and bandwidth efficiency on AMD XDNA NPUs. Empirical measurements for representative models (ViT, CLIP, BERT) on a Ryzen™ AI 9 HX (4×8 NPU grid, 60 GB/s DRAM) indicate the following:

Model	Unfolded Attention (ms)	Folded Attention (ms)	Speed-up
ViT-base-patch-16	3.2	0.8	4.0×
CLIP-patch32-MHA1	2.0	0.6	3.3×
CLIP-patch16-MHA1	5.0	1.5	3.3×
BERT	1.5	1.38	1.1×

End-to-end network latency improves by up to 32% for memory-bound models, with bandwidth utilization reduced by ∼8% even in compute-bound settings. Maximal per-attention speedup observed is 4× compared to naive, unfolded execution. Attention models with high ratio of DRAM movement-to-compute benefit most. For BERT, where attention is not the bottleneck, the reduction is smaller (Deshmukh et al., 25 Aug 2025).

4. KamLAND-Zen: Interpretable Temporal Attention in Rare Event Detection

Within KamNet, Zen-Attention provides interpretability and enhanced discrimination between rare signal and background in kiloton-scale liquid scintillator detectors (Li et al., 2022).

4.1 Role and Architecture

KamNet processes event data as a time series of sparse, spherical PMT hit maps (28 × 38 × 38 tensor for 28 time-slices and spatial pixels). Zen-Attention is inserted between stacked ConvLSTM layers (for temporal feature coding) and spherical CNNs (to exploit SO(3) symmetry). It computes normalized temporal scores $S(t)$:

$$S(t) = \mathrm{Softmax}_t \left[ \sum_{c,\theta,\phi} H_{c,t,\theta,\phi} \cdot W_{c,t,\theta,\phi} \cdot O_{c,\theta,\phi} \right]$$

and forms context images:

$$I_{\mathrm{context}}(c,\theta,\phi) = \sum_{t=1}^T S(t)\, H_{c,t,\theta,\phi}$$

Alternatively, Zen-Attention is expressible as scaled dot-product attention with queries from the final ConvLSTM output and keys/values derived from hidden states.

4.2 Quantitative Performance and Interpretability

Zen-Attention empowers KamNet to focus on discriminative temporal intervals (e.g., 5–10 ns for γ-cascade backgrounds, 0–5 ns for positronium delays), as revealed by projecting $K$0 over time. The resulting improvement in background rejection—at fixed 0νββ signal acceptance—is significant:

Model Variant	¹⁰C Rejection @ 90% 0νββ Acceptance
Planar CNN	61.5%
SO(3)+ConvLSTM, no attention	~67%
KamNet + Zen-Attention	74.0%

In full simulation of KamLAND–Zen 800, this yields ≳10% absolute increase in half-life sensitivity for the rare signal. Attention weights also clarify the most informative hardware time windows and guide detector design choices.

5. Limitations and Extensions

5.1 Compiler Framework Constraints

• Maximum context length is limited by aggregate L1 buffer capacity; very long sequences (N > 1024) cannot achieve full folding, reducing speedup benefit.
• Static sequence lengths are assumed per batch; highly dynamic inputs may experience padding overhead.
• The current implementation is finely tuned for a specific NPU ISA; porting requires re-engineering.

Future extensions discussed include runtime adaptive tiling, support for attention in training (including backward pass), integration via MLIR or LLVM for broader backend support, and extension to sparse or LSH-based attention patterns (Deshmukh et al., 25 Aug 2025).

5.2 Application-Specific Insights in KamNet

• The improvement from attention saturates for time binning finer than 1.5 ns; further subdivisions yield diminishing returns.
• Hardware upgrades targeting the early time window (e.g., LAPPDs, higher QE) and chemical modifications to scintillator timing profile are directly motivated by interpretability from Zen-Attention weights.
• For solar neutrino elastic-scatter signals with no temporal discriminant, Zen-Attention weights naturally become uniform, confirming model reliability.

6. Broader Significance

Zen-Attention as a compiler framework substantiates that coalescing entire transformer attention blocks into a single, spatially folded NPU operator can deliver order-of-magnitude improvements in DRAM efficiency and online inference latency—critical for deployment on edge and energy-constrained hardware. The interpretability-driven Zen-Attention module in KamNet demonstrates that temporal attention is not only an optimization but also an analysis tool, providing both enhanced statistical discrimination and principled detector design feedback (Deshmukh et al., 25 Aug 2025, Li et al., 2022).