Learning How to Remember: Memory Modality Selection for GUI Agents

AI summary

Overview

Research area: Artificial intelligence / vision-language models applied to GUI (graphical user interface) agents, specifically memory augmentation for mobile Android agents.

Technical level: Advanced. The paper assumes familiarity with vision-language models, transformer hidden states, probing classifiers, PCA, and retrieval-augmented agents.

Scope: This paper studies whether the memory given to a GUI agent should be text or an image, shows that the answer changes from step to step, and proposes a router (SwitchMem) that picks one modality per step.

What This Paper Is About

Memory-augmented GUI agents retrieve past experience to help decide the next action, but existing systems fix that experience in one form — always text, always screenshots, or always both. This paper asks whether a single fixed modality is actually the best choice at every step of an interaction. The authors show that the preferred modality shifts across steps, and that this preference can be read out of the agent's own hidden representations, which lets them build a lightweight router that selects text or image memory per step.

Key Contributions

1. A systematic study of memory modality in VLM-based GUI agents, showing that the preferred memory representation is step-dependent and that fixed modalities leave performance on the table across evaluation settings.
2. Directional Chain-of-Embedding (D-CoE), an analysis method that computes the signed difference between the actor's hidden states under Image and Text memory at each layer, revealing a reproducible Text/Image preference signal that generalizes to unseen GUI trajectories.
3. SwitchMem, a routing framework that uses this representation-level signal to select Text or Image memory at each interaction step, providing only the selected modality to the frozen actor.
4. Empirical validation across three GUI benchmarks (AITW, LearnGUI, MobileMiniWoB++) plus generalization to a GUI-specialized actor (UI-TARS-7B-SFT), with routing-behavior, ablation, budget-matching, and latency analyses.

Main Findings

• Neither modality dominates. Across AITW and LearnGUI, both Text and Image memory are preferred for substantial subsets of interaction steps. Image memory yields a slightly higher average performance on both, but the weaker modality on aggregate still provides higher utility on many individual steps.
• A step-wise oracle beats any fixed modality. Running the frozen actor separately with Text and Image memory under the same evaluation seed and taking the higher score consistently outperforms the strongest fixed-modality strategy, revealing a persistent gap caused by static memory representation.
• Modality preference is visible in hidden states. With the D-CoE shift, a PCA + linear probe predicts the preferred modality; AUROC increases in deeper layers. The probe retains only examples with at least five-of-seven agreement across sampled generations to reduce sampling noise.
• Best probe layers differ by benchmark. The modality-preference signal reaches its highest AUROC at layer 31 on AITW and at layer 23 on LearnGUI. SwitchMem nonetheless always uses layer 31 (the final transformer block) to keep one benchmark-independent extraction rule.
• SwitchMem leads on all three benchmarks. With the frozen Qwen3.5-9B actor: AITW 51.420%, LearnGUI 38.131%, MobileMiniWoB++ 63.043%, versus the strongest fixed memory conditions of 49.300 (Image), 37.441 (Image), and 60.870 (Both) respectively.
• Advanced memory formulations did not beat raw context injection. RAG-GUI scored 18.300 / 24.186 / 9.783; M3A 17.200 / 20.899 / 41.304; AppAgent 14.780 / 31.104 / 25.000; EchoTrail 45.900 / 27.286 / 55.435; AGMem 45.380 / 28.366 / 52.174.
• Routing is state-adaptive, not task-locked. SwitchMem uses both modalities within a single trajectory (83.8% dual-modality trajectories on AITW, 90.8% on LearnGUI, 57.8% on MobileMiniWoB++), with average switch counts of 1.72, 1.79, and 1.56 and contiguous run lengths of 2.42, 4.95, and 4.50.
• Image memory helps under identifiable conditions. Image-over-Text utility advantage is positively associated with element density (unique UI elements divided by their fractional screen coverage), is larger for Function-oriented actions than Spatial actions, and is larger in post-initial states than initial states. SwitchMem's router follows these same patterns.
• Actor-derived targets matter. Simpler routing baselines (Utility, Binary, Basic) offer minimal gains over a static Image policy on LearnGUI; Binary is closely competitive on AITW.
• Gains come from step-specific assignment, not the overall ratio. SwitchMem outperforms ratio-matched random routing at every evaluated intermediate Image-selection rate. Random routing's expected score is modeled as S_rand(p) = (1 − p)·S_Text + p·S_Image.
• Generalizes to another actor. With UI-TARS-7B-SFT: no memory 29.980 / 25.705 / 23.913; Text 32.640 / 38.230 / 32.609; Image 30.440 / 40.997 / 33.696; SwitchMem 33.680 / 41.055 / 36.957.
• Modest latency overhead. SwitchMem adds only 0.128 s, 0.190 s, and 0.304 s per step over the Image baseline on AITW, LearnGUI, and MobileMiniWoB++. RAG-GUI requires 8.815 s and 11.118 s per step on AITW and LearnGUI, versus 1.407 s and 1.848 s for SwitchMem (RTX PRO 6000).

Methodology in Plain English

The study holds everything constant except how the retrieved memory is presented. A single text-based retriever picks one past demonstration step for the current GUI state, then that same demonstration is instantiated either as text (instruction, textual state description, demonstrated action) or as an image (instruction, screenshot, demonstrated action). This isolates representation from retrieval quality.

To estimate how good each modality is at a given step, the frozen actor samples K actions per modality and scores them with the benchmark-specific action scorer against the ground-truth action. The difference between the Image and Text utilities defines the step-wise preference. A step-wise oracle then takes the better of the two scores at each step.

For the representation-level analysis, the authors extract the actor's hidden state at the final prompt token right before action generation, under both memory conditions, and take the signed difference (Image minus Text) at each layer — the Directional Chain-of-Embedding. After standardizing and reducing this 4096-dimensional shift to 64 dimensions with PCA, a logistic probe predicts which modality the actor prefers.

SwitchMem itself trains in two stages. First, a preference estimator fits a logistic probe on the final-layer shift (layer 31) to produce a continuous soft preference target t_i; ties are excluded. Second, a pairwise router is trained: a shared MLP scorer maps each (current state, memory candidate) pair to a scalar, and the router score is the Image score minus the Text score. The router is trained by MSE against the soft target. At inference the router score is compared with a threshold; Image is chosen if the score meets the threshold, otherwise Text, and only that one memory goes to the frozen actor.

Setup details: Qwen3.5-9B is the frozen actor and Qwen3-VL-Embedding-2B handles top-1 retrieval and memory embedding. Training samples seven actions per modality at temperature 0.7 with seeds 42–48. The router is a two-layer MLP (1024 → 256 → 1) with GELU activation and 0.1 dropout, trained with AdamW at a learning rate of 0.001. AITW uses 20,000 training steps split at trajectory level into 10,019 preference-query steps (3,100 trajectories) and 9,981 memory steps (3,087 trajectories) with no overlap; testing uses 5,000 steps from 759 held-out trajectories. LearnGUI uses 101 successful Online demonstration trajectories totaling 1,500 interaction steps. MobileMiniWoB++ provides no step-level utilities, so the LearnGUI router is transferred without weight updates; its memory bank has 129 interaction steps from 50 successful trajectories (suite seed 1001), the threshold is calibrated on initial states of 92 task templates (seed 43), and evaluation runs on the same 92 templates with seed 42, giving 63.043% task success (58/92).

Why This Matters

Impact on research. The paper reframes memory modality from a fixed architectural choice into a per-step decision, and backs that reframing with an interpretable internal signal (D-CoE). It also supplies a negative result of practical value: more elaborate memory formulations (procedural guidelines, operation documents, action-region crops) did not consistently beat simply injecting retrieved context, suggesting that how experience is represented at each step may matter more than how elaborately it is constructed.

Real-world applications.

• Mobile assistants that tap, swipe, and type on a user's behalf and must recover from mid-task errors by reusing past successful trajectories.
• Automated app testing and QA, where an agent replays and verifies multi-step flows across an app's screens.
• Accessibility tooling that drives an interface on behalf of users who cannot operate it directly.
• Enterprise process automation for repetitive cross-app workflows (form entry, navigation, status checks) on Android devices.

Industry relevance. The approach is training-free at the actor level (the actor stays frozen), adds only a small MLP router and a per-step latency of roughly 0.13–0.30 s over an image-memory baseline, and transfers to a GUI-specialized actor. That makes it attractive as a drop-in layer on top of an existing deployed agent rather than a full retraining effort.

Future Directions

• Beyond Android and emulation. The authors state that generalization to iOS interfaces and newer Android versions is unverified, and that online evaluation used emulated rather than physical devices, leaving device rendering differences, system interruptions, network variability, and hardware constraints untested.
• Routing without hidden-state access. SwitchMem requires the actor's hidden states both for preference target construction and for inference-time routing, so it cannot be applied directly to models served through APIs that do not expose those representations. Extending modality routing to such models is stated as an open direction.
• Broader modality and memory design space. The current framework chooses only between Text and Image for a single retrieved demonstration; whether richer modality sets, multi-memory contexts, or non-GUI visual environments behave the same way is not addressed.
• Benchmark coverage for probing. A 5-of-7 preference probe was not reported for MobileMiniWoB++ because it lacks intermediate step-level utilities, so adding such signals to more benchmarks could test whether the D-CoE signal holds there directly rather than only through transfer.

Target Audience

Researchers and engineers working on VLM-based GUI agents, agent memory architectures, and retrieval-augmented multimodal systems. It is most useful to readers comfortable with transformer internals and probing methods, and to practitioners building mobile automation who want an adaptive memory layer that can sit on top of a frozen model. Readers looking for an introductory treatment of GUI agents would find the probing and routing formulation demanding.

Authors’ abstract

Recent advances in vision-language models (VLMs) have substantially improved the capabilities of GUI agents, while memory augmentation has emerged as a promising approach for reusing knowledge from past interactions. However, existing memory augmented GUI agents typically adopt a fixed representation strategy, providing retrieved experience as a static modality. It remains unclear whether a single memory modality is consistently optimal, or how modality preference changes across interaction steps. In this work, we conduct a systematic study of memory modality in VLM based GUI agents. Across GUI benchmarks, we find that the preferred memory representation is strongly step dependent. Motivated by the finding, we introduce \textbf{SwitchMem}, a routing framework that dynamically selects between textual and visual memory representations at each step. Across three GUI benchmarks, SwitchMem outperforms memory augmented baselines that rely on static memory representations. Our findings validate our claim that memory modality is a step dependent decision rather than a fixed design choice. Code and project resources are available at https://txxnrd.github.io/switchmem