Probabilistically Guaranteed Retention Architecture

• Probabilistically Guaranteed Retention Architecture is an adaptive mechanism that learns per-token retention probabilities under a strict global budget, ensuring a deterministic top-M selection at inference.
• It employs a lightweight retention module with Hard-Concrete relaxation and dual optimization to efficiently reduce the quadratic attention cost without altering core model components.
• Empirical studies show that retaining 30%-50% of tokens preserves nearly full model performance while significantly reducing memory usage and improving throughput.

Searching arXiv for the cited papers to ground the article in current literature. {"query":"arXiv (Rafiuddin et al., 9 Oct 2025) Learning What to Remember Adaptive Probabilistic Memory Retention for Memory-Efficient LLMs", "max_results": 5} {"query":"(Rafiuddin et al., 9 Oct 2025)", "max_results": 10} Probabilistically Guaranteed Retention Architecture denotes an adaptive probabilistic memory-retention mechanism for Transformer encoders in which each token representation is assigned a learned retention probability under a strict global budget $M$, and deployment converts the learned stochastic policy into a deterministic top-$M$ selector that retains exactly $M$ tokens. In "Learning What to Remember: Adaptive Probabilistic Memory Retention for Memory-Efficient LLMs" (Rafiuddin et al., 9 Oct 2025), the term refers specifically to an architecture-agnostic, layer-wise token selection module that reduces the active token set passed from one layer to the next without modifying the base attention mechanism or the downstream task head. The motivation is the quadratic scaling of Transformer attention, $O(n^2)$, which limits long-context use; the proposed retention mechanism addresses this by learning what to keep rather than processing and storing all token states.

1. Formal problem setting

The architecture is defined over an input sequence

$\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$

with encoder hidden states

$$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$

Its objective is to learn which token states should be retained under a strict budget $M$, rather than keeping all $T$ tokens. Retention is represented by binary variables

$$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$

yielding masked hidden states $\mathbf{H}\odot\mathbf{z}$.

The retention policy is learned as a constrained stochastic optimization problem: $M$0 where $M$1 is the learned retention probability and $M$2 is the maximum expected number of retained tokens. The defining property is therefore not heuristic pruning, but learning a distribution over retention decisions under a global budget (Rafiuddin et al., 9 Oct 2025).

This formulation makes the retained set an endogenous part of model optimization. The budget is applied to the expected number of retained tokens during training, and the encoder is trained jointly with the retention policy. A plausible implication is that the architecture treats memory as a first-class optimization variable rather than as an after-the-fact compression step.

2. Retention module and layer-wise scoring

The retention module is lightweight and plug-in compatible with standard encoders. At each layer it maintains a summary state

$M$3

and computes a retention score

$M$4

followed by

$M$5

The decision depends on three stated factors: local token content $M$6, global history via $M$7, and a learned sigmoid probability $M$8. The decay parameter $M$9 controls how much prior context is remembered in the scorer. Because the module only inserts a learned retention mask that shrinks the active token set across layers, the base encoder remains intact; the method is described as a drop-in efficiency module for standard Transformer encoders such as DistilBERT or Longformer (Rafiuddin et al., 9 Oct 2025).

This design is significant because it separates retention from attention itself. The architecture does not alter the attention equations, does not require specialized sparse-attention kernels, and does not require changing downstream task heads. Its operational effect is instead to reduce the number of token states forwarded between blocks.

3. Constrained learning, dual optimization, and relaxation

To enforce the budget, the paper introduces a Lagrange multiplier $M$0: $M$1 Training is posed as the saddle-point problem

$M$2

Optimization alternates between SGD on $M$3 and projected gradient ascent on the dual variable: $M$4 The stated interpretation is that budget pressure becomes adaptive: if too many tokens are being retained, $M$5 increases and pushes the probabilities down.

Because $M$6 is discrete, direct backpropagation is not possible. The paper therefore uses a Hard-Concrete reparameterization. With

$M$7

the relaxed gate is

$M$8

where $M$9 is a temperature and the stretch parameters satisfy $O(n^2)$0. During training, $O(n^2)$1 is replaced by $O(n^2)$2 in the loss and Lagrangian. The stated effects are low-variance gradients, a differentiable approximation to Bernoulli sampling, and approximate enforcement of sparsity and budget constraints. The paper also characterizes this as a variational relaxation of discrete token gating (Rafiuddin et al., 9 Oct 2025).

The combination of dual optimization and Hard-Concrete gating is the core probabilistic machinery of the architecture. The learned object is a distribution over retention masks, but the distribution is shaped by an explicit resource constraint rather than an unconstrained salience objective.

4. Inference-time enforcement and the meaning of the guarantee

At test time, the probabilistic model is converted into a hard retention mask using the top-$O(n^2)$3 probabilities. If $O(n^2)$4 is the $O(n^2)$5-th largest value in $O(n^2)$6, then

$O(n^2)$7

so that

$O(n^2)$8

The method therefore retains exactly the $O(n^2)$9 highest-probability tokens (Rafiuddin et al., 9 Oct 2025).

This top-$\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$0 rule is the practical source of the “probabilistically guaranteed” designation. During training, the model learns probabilities of importance and constrains expected retention; during inference, selection becomes deterministic and budgeted. Memory usage is therefore controlled by design rather than by post hoc truncation.

A common misconception is to interpret the guarantee as claiming that every stochastic sample during training exactly obeys the budget. The paper explicitly does not make that claim. Its interpretation is narrower and more precise: the expected retention is budgeted during training, the hard inference mask respects the budget at deployment, and dual optimization encourages the learned probabilities to align with the constraint. The appendix is described as giving a slackness-style interpretation in which the dual variable and the duality gap bound how far the learned probabilities can deviate from the target budget.

The guarantee is thus hybrid. It is probabilistic at training time, because the policy is expressed as Bernoulli retention variables under an expectation constraint, and deterministic at deployment, because the top-$\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$1 rule enforces an exact operational memory bound.

5. Empirical behavior, efficiency, and ablations

The reported experiments cover six benchmarks across classification, extractive QA, and summarization: SST-2, IMDb, CUAD, QASPER, ArXiv, and PubMed RCT. Backbones are DistilBERT-base-uncased for SST-2, IMDb, and CUAD, and Longformer-base-4096 for ArXiv, QASPER, and PubMed. Budgets tested are $\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$2 token retention and $\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$3 token retention (Rafiuddin et al., 9 Oct 2025).

Setting	Benchmarks	Backbone
Classification	SST-2, IMDb, CUAD	DistilBERT-base-uncased
Extractive QA / Summarization	QASPER, ArXiv, PubMed RCT	Longformer-base-4096

Across tasks, keeping only $\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$4 to $\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$5 of tokens preserves $\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$6 of full-model performance while cutting peak memory by $\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$7 to $\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$8 and improving throughput by up to $\mathbf{X}=(\mathbf{x}_1,\dots,\mathbf{x}_T),$9. At $$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$0 retention, results are typically within about $$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$1–$$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$2 percentage points on many tasks. At $$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$3 retention, performance remains strong, usually within about $$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$4–$$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$5 percentage points on short classification tasks and similarly tight on long-document tasks.

Concrete examples reported in the paper include SST-2 at $$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$6 versus $$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$7 at $$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$8 retention and $$\mathbf{H}=(\mathbf{h}_1,\dots,\mathbf{h}_T), \qquad \mathbf{h}_t\in\mathbb{R}^d.$$9 versus $M$0 at $M$1; IMDb at $M$2 versus $M$3 at $M$4 and $M$5 versus $M$6 at $M$7; ArXiv $M$8 at $M$9 versus $T$0 at $T$1 and $T$2 versus $T$3 at $T$4; and QASPER F1 matching dense F1 at both budgets, $T$5 and $T$6. CUAD is described as staying very close to the full model on micro/macro-F1, and PubMed RCT as showing small ROUGE drops, with ROUGE-L matching the dense model in the reported results.

The comparison set includes random pruning, H2O fixed pruning, constraint-aware pruning, and Infor-Coef; the method is reported to outperform these baselines and to be competitive with or slightly better than sparse-attention baselines such as Longformer and BigBird on the tested tasks. Efficiency numbers in the compute table include $T$7 s/batch and throughput $T$8 for the full Transformer, $T$9 s/batch and throughput $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$0 for Adaptive Retention at $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$1, $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$2 for Longformer at $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$3, and $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$4 for BigBird at $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$5. At $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$6 retention, the paper reports $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$7 throughput and a $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$8 GB memory footprint on a $$z_t = \begin{cases} 1, & \text{if } \mathbf{h}_t \text{ is retained} \ 0, & \text{otherwise,} \end{cases}$$9 GB GPU.

Ablation results indicate that each major component contributes materially. Removing Hard-Concrete relaxation hurts accuracy and reduces throughput; disabling alternating optimization hurts; fixing $\mathbf{H}\odot\mathbf{z}$0 weakens constraint handling; and threshold-based pruning performs worst among the non-random variants. The reported interpretation is that the probabilistic plus constrained-optimization formulation is not merely stylistic but functionally important.

6. Relation to other memory-retention architectures

The term “retention architecture” also appears in other parts of the literature, but with different formal commitments. "ATLAS: Universal Function Approximator for Memory Retention" (Deventer et al., 2022) addresses continual learning rather than Transformer context compression. ATLAS is presented as an architecture with intrinsic memory retention based on parameter isolation, locality, sparse support, and distal orthogonality. Its relevant formal properties are deterministic: sparsity, bounded gradient norm, and the distal orthogonality condition

$\mathbf{H}\odot\mathbf{z}$1

The paper explicitly states that ATLAS does not provide a probabilistic retention guarantee of the form “with probability at least $\mathbf{H}\odot\mathbf{z}$2, memory is preserved.” Its retention claim is therefore deterministic in structure and empirical in effect, not probabilistic in the formal sense.

A different adjacent formulation appears in "Learning What to Remember: Observability-Safe Memory Retention via Constrained Optimization for Long-Horizon Language Agents" (Kang et al., 9 Jun 2026). There, memory retention is framed as a constrained stochastic optimization problem over a retained subset

$\mathbf{H}\odot\mathbf{z}$3

with expected cumulative reward

$\mathbf{H}\odot\mathbf{z}$4

OSL-MR is described as observability-safe because it strictly separates online-observable inputs from offline-available supervision, and its deployed policy uses a budget-constrained decoder to ensure feasibility. This is a probabilistically grounded retention architecture for long-horizon language agents, but it targets evidence utility, miss penalties, reacquisition costs, and stale-information risk under realistic observability constraints rather than token-state retention inside an encoder.

These comparisons clarify the scope of the phrase. In (Rafiuddin et al., 9 Oct 2025), a Probabilistically Guaranteed Retention Architecture is specifically a Transformer-side mechanism that learns per-token retention probabilities under a global budget and enforces an exact top-$\mathbf{H}\odot\mathbf{z}$5 budget at inference. In (Deventer et al., 2022), memory retention is structural and deterministic. In (Kang et al., 9 Jun 2026), retention is probabilistically grounded at the level of sequential decision-making for agent memory. This suggests that the distinctive contribution of the encoder-focused formulation is the combination of probabilistic token selection during training with deterministic budget enforcement at deployment, yielding a practical memory guarantee without redesigning the Transformer itself.