PaperlayerView exact revision on arXiv

Source-linked AI summary

AURA: Action-Gated Memory for Robot Policies at Constant VRAM

Josef Chen

arXiv:2606.02775v1cs.AIcs.ARcs.DCcs.PFcs.RO

TL;DR

Robots running long, non-resetting episodes face growing KV-cache memory and write-bandwidth costs that differ from datacenter inference. AURA-Mem replaces that cache with constant-size recurrent memory and a learned gate trained on closed-loop action error. It reports accuracy parity with O(1) baselines at substantially fewer writes and constant state, while its instantiated value-loss bound is vacuous at current scale.

  • Problem

    Long-horizon robots need memory under bandwidth and finite-write constraints, but KV-caches grow with episode length and generic compression objectives may retain information unrelated to action utility.

  • Method

    AURA-Mem wraps a policy with fixed-size recurrent fast-weight memory and an action-error write gate trained using a closed-loop action-information-bottleneck objective.

  • Results

    4.98–9.19× fewer writes per second are achieved at statistically equivalent synthetic-task accuracy, with 4,224 bytes of constant state versus a 6,061× larger growing-KV reference at 100,000 steps.

  • Takeaways & Limitations

    The results establish a write-bandwidth frontier in which action-utility gating preserves supported task accuracy while reducing memory traffic and keeping inference state O(1).

  • Takeaways & Limitations

    The primary N=64 write-bandwidth result uses n=3 seeds on a saturated task, so bandwidth efficiency rather than an accuracy gap is the meaningful discriminator.

Abstract

from arXiv · show

The KV-cache is the right memory for datacenters but the wrong memory for robots. Datacenter inference batches many short requests and resets them, amortizing an attention cache across a crowd. Embodied agents instead run one long, non-resetting episode on bandwidth-limited edge hardware, where high-bandwidth memory and flash are scarce, flash has finite write endurance, and memory writes rather than compute can become the binding constraint. AURA-Mem (Action-Utility Recurrent Adaptive Memory) targets this regime. It wraps a frozen vision-language-action backbone with a constant-size recurrent memory and a learned gate that writes only when the current observation would change the next action: memory that knows when to stay silent. Unlike reconstruction-based memory, the gate is trained directly against a closed-loop action-error signal. Its inference state is fixed at 4,224 bytes regardless of horizon, while a KV-cache grows to 6,061 times larger at 100,000 steps. On a controlled synthetic benchmark, AURA-Mem matches the best O(1) baseline in accuracy while using 5.19-6.13 times fewer writes, and up to 9.19 times fewer writes on easier configurations. Budget-matched random and periodic schedules do not recover this gain, isolating the benefit to the action-surprise signal. On a trained closed-loop OpenVLA-OFT 7B panel on LIBERO-Long (n=60 episodes per arm), the gate does not hurt success: AURA-Mem matches the ungated base policy (0.233) and slightly exceeds an always-write KV arm (0.217), while using 7.0 times fewer writes and constant memory. We also instantiate an approximate-information-state value-loss bound as a methodology demonstration; at this scale, the bound is vacuous rather than a guarantee.

1 Introduction

AURA-Mem targets robots’ unbounded, bandwidth-limited episodes by replacing growing KV-caches with constant-size, action-sufficient memory that writes selectively. Its reported gains are lower write traffic and fixed inference-state VRAM, while its action-sufficiency bound remains a methodology demonstration rather than a guarantee.

  • Deployment regime: The design addresses robots’ batch-1, continuous episodes, where growing KV-caches cannot amortize cost across requests and memory bandwidth can dominate inference latency.Datacenter inference instead batches short requests and resets sessions, making growing KV-caches a different regime.
  • Action-gated memory: AURA-Mem uses a learned action-utility write gate trained through closed-loop action loss rather than reconstruction or token-level objectives.The gate is driven by the policy’s action-prediction error and is paired with an action-information-bottleneck objective.
  • Write-bandwidth frontier: 4.98–9.19× fewer memory writes per second than a dense baseline at statistically equivalent accuracy on memory-dependent synthetic benchmarks.Budget-matched random and periodic schedules do not recover the accuracy at the same write rate.
  • Constant inference-state VRAM: 4,224 bytes of fixed inference-state memory remains constant across 100,000 steps, while a matched growing-KV reference reaches 25,600,000 bytes.The reference is 6,061× larger at 100,000 steps; the comparison uses an analytic extrapolation for the long-horizon KV figure.
  • Action-sufficiency bound: The instantiated action-sufficiency value-loss bound is vacuous at current scale and is reported as a methodology demonstration, not a formal guarantee.Measured action-prediction sufficiency is strong, but the loaded bound exceeds the trivial value span.
  • Scope of evaluation: AURA-Mem’s real-VLA panel is explicitly framed as a zero-shot proof-of-mechanism rather than a state-of-the-art success sweep.The memory and measurement layer does not by itself raise robot success.

2 Related Work

Related work provides bounded recurrent or compressed memory, but AURA-Mem combines action-utility gating, closed-loop training, write-rate control, constant-state memory, and an empirical action-sufficiency certificate. Its closest alternatives generally retain growing caches, use surrogate objectives, or lack formal sufficiency measurements.

  • Bounded recurrent memory: AURA-Mem differs from fixed-size recurrent and linear-attention methods through action-utility surprise gating, closed-loop action training, and a controllable write-rate frontier.The gate is not based on a step clock, perplexity gradient, or spatial prediction.
  • KV-cache compression and eviction: KV eviction methods reduce cache footprint but remain asymptotically growing with context length rather than guaranteeing O(1) VRAM.They bound the cache to a window or fraction of history, not to a fixed-size state.
  • Embodied memory: VLA-specific compression methods use token-level surrogate objectives or eviction gates, whereas AURA-Mem trains against a closed-loop action loss and provides an action-sufficiency certificate.The comparison applies to VLA-Cache and KV-Efficient VLA in the supplied related-work discussion.
  • RL and robot-policy memory: Classical RL memory methods establish the need for bounded recurrent state but do not connect state size to a formal sufficiency bound.Recent robot-policy memory works use growing external stores or unconditional write schedules and provide no certificate in the cited comparison.
  • Surprise signals: AURA-Mem repurposes forward-model prediction error from intrinsic exploration into a memory-management signal for deciding when to update fast-weight state.In the cited comparison, exploration methods use surprise for action selection rather than memory management.
  • State representation theory: The AIS framework supplies the theoretical backbone, but AURA-Mem’s measured ε values are informative while its value-loss bound is vacuous at current scale.The certificate concerns value sub-optimality from finite memory, distinct from trajectory generalization certificates.

3 Method: AURA-Mem

AURA-Mem wraps a frozen VLA backbone with one bounded fast-weight state and a learned action-utility gate. The state is read every tick but updated only when incoming information is sufficiently novel for action selection.

  • Architecture: AURA-Mem carries one fixed-size recurrent fast-weight state across the episode instead of evicting tokens from a growing KV-cache.The state contents evolve on selected steps, while its shape remains independent of horizon.
  • Architecture: The frozen OpenVLA-OFT 7B backbone summarizes observations into tokens, while only lightweight projections, gate, and memory heads are trained.The backbone weights remain unchanged throughout the method.
  • Gated memory updates: The fast-weight state reads a query through W before writing, producing a memory output that conditions the action head.Keys and values support an associative fast-weight map, while the read is a linear query-to-state operation.
  • Memory and trade-offs: 4,224 bytes is the measured inference-state footprint for the batch-1, fp32, dk=dv=32 configuration, independent of horizon.At 100,000 steps, the matched KV reference reaches 25,600,000 bytes, or 6,061× larger; training-time activations remain O(T).
  • Gated memory updates: Reading occurs every control tick, but the learned binary gate skips writes when the current step lacks action-relevant novelty.When the gate is zero, the state is carried over unchanged and no write traffic is incurred.
  • Training objective: The state is trained against closed-loop action prediction through an action-information bottleneck rather than generic reconstruction alone.The action loss pressures the memory to retain information relevant to control, while the information-bottleneck term acts as a compression objective rather than a certificate.
  • Memory and trade-offs: The method traces a success–bytes–writes/sec curve by holding the target write rate fixed while sweeping state size.The paper does not claim a full two-dimensional Pareto sweep over state bytes and write rate.

4 Theory: an action-sufficiency value-loss bound

The paper instantiates an existing approximate-information-state framework for the recurrent memory and connects reward/action and transition prediction errors to value loss. Its empirical readout is diagnostic rather than a formal deployment guarantee.

  • AIS construction: AURA-Mem’s recurrent fast-weight state serves as the fixed-size compressed state in an approximate-information-state construction.The framework attaches learned reward/action and latent-transition heads to the compressed state.
  • AIS construction: The bound requires bounded rewards, an integral probability metric, finite value-function regularity, time-homogeneous models, and coverage over reachable histories.These are the stated regularity and reachability assumptions for the instantiated theorem.
  • Bound structure: The action-sufficiency bound combines reward/action error ε and transition error δ as η = ε + γLVδ, then scales the resulting approximation by 1/(1−γ).The contraction step yields α ≤ η/(1−γ), followed by a factor-of-two policy-loss argument.
  • Bound structure: The tighter version replaces LVδ with the measured one-step value-prediction residual Δ*(V̂*) when that residual is available.The paper reports this form as never looser than the LVδ form.
  • Empirical certification: The empirical certification estimates ε and δ with Monte Carlo diagnostics on visited query steps rather than verifying worst-case sup-errors over all reachable sequences.The reported quantities are action-prediction and policy-versus-teacher next-state/action distribution diagnostics.
  • Empirical certification: The instantiated bound is numerically vacuous at the current scale, so the paper presents it as a methodology demonstration rather than a formal guarantee.The authors explicitly disclaim a novel theorem, an exact sufficient statistic, and an off-distribution guarantee.

5 Experiments

The experiments evaluate write bandwidth, accuracy, and constant inference memory on synthetic memory-stress tasks, with comparisons spanning gated, ungated, random, periodic, token-loss, recurrent, and no-memory variants.

  • Evaluation scope: The study uses synthetic tasks without real-world data or external-corpus pretraining, with up to seven seeds per cell and disclosed smaller-sample exceptions.Accuracy and bandwidth sweeps use H100 GPUs; the constant-VRAM measurement uses a real L40S GPU.
  • Tasks: The primary benchmark is NoisyLongRecall, which tests multi-binding associative recall under distractors and overwrite operations.SparseRecall separately tests rare event tokens and illustrates gate selectivity and the write-rate knob.
  • Main findings: 4.98–9.19× fewer memory writes is achieved than the write-every-step dense baseline while maintaining statistically equivalent accuracy on matched synthetic benchmarks.The paper identifies the write-bandwidth frontier, rather than accuracy superiority, as the main claim.
  • Main findings: All non-trivial variants reach approximately 1.0 accuracy on easy saturated configurations, making writes/sec and constant VRAM the differentiating axes.The main state-budget sweep uses N=dk=dv ∈ {16, 32, 64}, with matched parity variants at each budget.
  • Variants and metrics: AURA-Mem is compared with write-every-step, fixed-size-state, GRU recurrence, random, periodic, learned-token, and no-memory controls.The fixed-size-state arm isolates the contribution of the learned gate, while random and periodic schedules test whether rate alone explains performance.
  • Main findings: Random and periodic schedules fail to recover AURA-Mem’s accuracy at the same write rate, while the learned token-loss gate collapses to g=0 always in the reported control.These controls isolate the role of the action-surprise signal and action-utility training objective.
  • Variants and metrics: Success is masked argmax accuracy on query steps, while writes/sec equals expected gate activity multiplied by the fixed 20.0 Hz control rate.Memory bytes are measured directly for O(1) variants, alongside ε and δ diagnostics.

6 Results

Across synthetic and closed-loop evaluations, AURA-Mem preserves task success while reducing memory writes and keeping its carried inference state constant. The results support bandwidth and VRAM efficiency claims, but not improved robot success or a full Pareto characterization.

  • Write-bandwidth frontier: 9.19× fewer writes per second than a matched dense baseline at statistically identical task success on the primary benchmark.At N=64, AURA-Mem records 2.18 ± 0.55 writes/s versus 20.0 writes/s, with success 1.000 ± 0.000 versus 0.9995 ± 0.0015.
  • Gate mechanism: Budget-matched random and periodic schedules lose accuracy, isolating the benefit to the learned action-utility gate rather than write frequency alone.At N=64, random and periodic schedules score about 0.365–0.375 versus AURA-Mem’s 1.000; the learned token gate collapses to g=0.
  • Scope and limitations: The experiments establish bandwidth and VRAM advantages, but do not show improved robot success or profile per-step wall-clock latency.The write-rate target is held fixed, tracing a one-dimensional curve rather than a full two-dimensional Pareto sweep.
  • Constant inference-state VRAM: 4,224 bytes of carried inference state remain constant through 100,000 steps, versus 25,600,000 bytes for a matched growing-KV stub.The KV comparison is an analytic extrapolation for the matched-dimension stub; the fixed state shape is independent of horizon.
  • Trained KV-cache head-to-head: Trained KV and AURA-Mem reach accuracy parity, while the KV cache grows linearly and exceeds AURA-Mem’s state after the crossover horizon.At batch 1, the trained KV cache reaches 262,144 bytes at T=1,024 and 2.56 MB at T=10,000, compared with AURA-Mem’s constant 4,224 bytes.
  • Closed-loop policy evaluation: On LIBERO-Long, AURA-Mem matches the base policy at 0.233 success and slightly exceeds the always-write KV arm at 0.217, using 7.0× fewer writes.The panel uses n=60 episodes per arm; absolute success reflects the zero-shot, single-seed evaluation regime.
  • Action-information-state analysis: The action-information-state value-loss bound is vacuous at the tested scale, so it is reported as a methodology demonstration rather than a formal guarantee.The real VLA stream also preserves the constant 4,224-byte state property.

7 Limitations and honest disclosures

The paper qualifies its main claims by separating structural memory benefits from accuracy superiority and by disclosing scope, evaluation, and ablation limits. Several results rely on synthetic or single-seed settings, while the AIS certificate is explicitly vacuous at the tested scale.

  • Scope of O(1) memory: O(1) inference memory applies only to the carried state, not training: BPTT retains activations across the full unroll length T.The 4,224-byte state and the roughly 43.1 MB CUDA allocator peak are distinct quantities.
  • Claim calibration: AURA-Mem matches, rather than exceeds, the strongest O(1) fixed-size baseline in accuracy while reducing writes by 4.98–9.19×.The paper frames write-bandwidth efficiency, not accuracy improvement, as the primary contribution.
  • Claim calibration: The 9.19× headline comes from a near-saturated N=64 task, where ceiling performance limits evidence about accuracy robustness.The hard task was introduced to recover a non-saturated regime, where AURA-Mem still reports near-parity accuracy with fewer writes.
  • Capacity control: AURA-Mem has greater gradient-active capacity than scheduled-write comparators, so the observed frontier may reflect both gating and added gate-module capacity.The paper states that a capacity-controlled ablation has not been run.
  • Theoretical certificate: The action-information-state value-loss bound is vacuous at this scale and is presented as a methodology demonstration rather than a formal guarantee.Action-prediction premises are strong, but the guaranteed bounds remain numerically uninformative.
  • Ablation limits: The information-bottleneck term is only suggestively beneficial because the token-gate comparator collapsed and the isolated ablation was not statistically decisive.The reported β ablation shows a +0.175 gap, but Welch p=0.153 and the write-rate effect is not significant at this sample size.
  • Configuration disclosure: The released checkpoint and sweep use different hyperparameters, so reported experimental metrics do not directly characterize the released configuration.The sweep uses dk=dv=N, whereas the checkpoint uses dk=dv=64 and a different inference-state size.
  • Evaluation scope: The evaluation remains bounded to synthetic recall benchmarks, with no real-robot deployment, energy measurements, or wall-clock latency profile.The VLA panel is described as a single-seed mechanism illustration, and latency crossover requires hardware profiling.

8 Conclusion

AURA-Mem combines bounded fast-weight memory with an action-error write gate and closed-loop action-IB training. Its results show fewer writes and constant state size, while the information-state bound remains vacuous and real-robot validation is absent.

  • AURA-Mem combines one bounded state object, a learned action-error write gate, and a closed-loop action-IB objective for O(1) inference-state VRAM.
  • 4.98–9.19× fewer writes per second are achieved than the dense baseline at statistically equivalent accuracy, while random and periodic schedules fail at the same write rate.The write-bandwidth result spans N=32–64 on the noisy long recall task.
  • 4,224 bytes of state remain flat across 100,000 steps, versus a 6,061×-larger growing-KV reference.
  • εmean=0.0021 accompanies an AIS value-loss bound that is vacuous at current scale, so the certificate is reported as a methodology demonstration rather than a formal guarantee.
  • The evaluation lacks real-robot deployment, energy measurements, and wall-clock latency profiling, while the IB term is not decisively isolated.
  • Future work targets physical-hardware deployment, trained KV-eviction comparisons, latency profiling, tighter certificates, and capacity-controlled gating ablations.

A.1 Proof of Theorem 4.3 (full detail)

The proof bounds value loss by comparing true and surrogate Bellman operators, propagating one-step mismatch through contraction, and accounting for policy-evaluation error. The resulting factor-of-two bound matches the cited theorem.

  • Notation: The true and surrogate Bellman operators are γ-contractions, with V∗ and ˆV∗ as their respective fixed points.
  • Step 1: dual (IPM) inequality: The IPM inequality bounds the surrogate value-function expectation mismatch by its Minkowski functional times the distance between true and surrogate next-state laws.
  • Step 2: one-step Bellman mismatch: The one-step Bellman mismatch is then bounded at every reachable history-action pair using the approximate-state assumptions.
  • Step 3: contraction propagates the one-step error: α = suph |V∗(h) − ˆV∗(σ(h))| propagates the one-step error through γ-contraction, yielding the approximate value discrepancy.
  • Step 4: value approximation to closed-loop loss: The final policy-value comparison adds a second error term bounded by α, producing the factor-of-two result for the greedy surrogate policy.The cited passage states that this factor matches Subramanian et al. Theorem 9/27.

A.2 What LV is, per metric

LV measures how the surrogate optimal value function scales the IPM transition discrepancy. Its form depends on the metric, and the instantiated bound is vacuous at the current experimental scale.

  • LV is the Minkowski functional of ˆV∗ with respect to the IPM function class F.
  • Total variation: For total variation, LV = 1/2 span(ˆV∗), bounded by half the reward range divided by 1−γ.
  • Wasserstein-1: For Wasserstein-1, LV equals the Lipschitz constant of the value function on the underlying metric space.
  • 52.69 versus a trivial value span of 10.0 makes the LV-loaded bound vacuous at current experimental scale.

B.2 Parameter-count disclosure

Parameter disclosures distinguish nominal architecture counts from gradient-active capacity and identify an unused predictor head. AURA-Mem has additional active gate parameters relative to scheduled-write variants.

  • Parameter totals scale with state dimension N=dk=dv, while the five primary parity variants have identical totals at each budget.
  • At N=32, AURA-Mem has 21,447 gradient-active parameters versus 15,110 for scheduled-write variants.
  • Predictor-head disclosure: The predictor sub-head contributes 8,320 nominal parameters but is never called and has no gradients in any variant.
  • +6,337 active parameters, or +41.9%, arise from the gradient-active gate MLP, creating a capacity advantage over scheduled-write variants.
  • GRU configuration: The corrected full-recurrence configuration uses GRU hidden=53; hidden=32 results were invalid and discarded.

B.3 Reproducibility commands

This section lists commands for deploying, running, aggregating, and checking AURA experiments, plus a reference training configuration and source location.

  • Experiment workflow: The deployment sweep is launched with aura/lean_sweep.py using Modal and up to 10 concurrent jobs.
  • Experiment workflow: Results are collected with aura.lean_aggregate using a run tag, then proof gates are checked with check_proof_matrix.py.
  • Reference configuration: The reference training run uses the ours variant on noisy_long_recall with seeds 0, 1, and 2 for 4,000 steps.
  • Reference configuration: The reference configuration specifies sequence length 96, batch size 64, learning rate 3e-3, write target rho 0.15, state dimension 32, and CUDA.
  • Release: Source code and the checkpoint are hosted at the Kaikaku/aura Hugging Face repository, which is currently private and available on request.
  • Extended measurement: Figure 14 extends Figure 7’s constant-versus-growing memory comparison to 100,000 steps and quantifies a 6,061× separation.

C Model card (abridged)

The model card describes intended research use, configuration distinctions, evaluation limitations, and broader deployment cautions. It also notes that the reported value-loss bound is vacuous at the current scale and that energy savings were not measured.

  • Intended use: AURA-Mem is intended for research benchmarking, ablations, and proof-of-concept long-horizon control on memory-constrained hardware.
  • Intended use: It is not validated for real-robot deployment, safety-critical or unmonitored use, or non-embodied workloads.
  • Configuration: The released checkpoint uses dk=dv=64 and a 16,640-byte inference state, unlike the dk=dv=N sweep configuration used for experimental results.
  • Limitations: All evaluation uses synthetic noisy-long-recall and sparse-recall benchmarks, with no real-robot validation.
  • Limitations: The ε and δ values are action self-consistency diagnostics rather than formal theorem-certificate values, and the instantiated value-loss bound is vacuous at current scale.
  • Energy considerations: The 4.98–9.19× write reduction is a proxy for possible energy effects; measured energy savings require hardware experiments beyond the paper’s scope.
  • Deployment caution: Because efficiency may accelerate deployment before validation, the model card calls for human oversight and task-specific safety validation.
Loading 2606.02775v1…