🤖 AI

Task-Aware Structured Memory for Dynamic Multi-modal In-Context Learning

The paper introduces TASM, a training-free framework for dynamic multi-modal in-context learning that overcomes context window and KV cache limitations by employing task-vector guided compression, semantics-aware token merging via bipartite graph matching, and a hierarchical memory structure to enable efficient, query-adaptive retrieval without disrupting semantic integrity.

Original authors: Zhirui Chen, Ziwei Chen, Ling Shao

Published 2026-06-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhirui Chen, Ziwei Chen, Ling Shao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very smart robot (a Multi-modal Large Language Model) how to do a new job by showing it thousands of examples. This is called "In-Context Learning." The robot looks at your examples and tries to guess the answer to your new question.

However, there's a big problem: The robot has a very short-term memory (a "context window"). If you show it too many examples, it runs out of space, gets overwhelmed, and slows down to a crawl. It's like trying to read a library of books in one sitting; eventually, you just can't hold all the information in your head.

Existing solutions try to fix this by being "ruthless editors." They look at the examples and delete the ones they think are boring or unimportant. But this paper argues that these editors are too clumsy. They often throw away crucial details, especially in pictures, breaking the spatial relationships (like where things are located) and leaving the robot confused.

Enter TASM (Task-Aware Structured Memory).

The authors propose a new, smarter way to manage the robot's memory. Instead of just deleting things, TASM organizes and compresses them. Here is how it works, using simple analogies:

1. The "Task Compass" (Task-Vector Guided Compression)

The Problem: Old methods decide what to keep based on what the specific example said. It's like a librarian who only keeps books that mention a specific word from a previous conversation, ignoring the rest of the story. This creates bias.
The TASM Solution: TASM first figures out the "direction" of the job. Imagine you are teaching a robot to identify cats. TASM creates a "Compass" that points toward "Cat-ness."

Instead of looking at individual examples, it looks at the transformation from a question to an answer.
It keeps the information that aligns with this "Compass" (the general logic of the task) and discards the noise that doesn't fit.
Result: The robot remembers the essence of the task, not just the specific details of the examples it saw.

2. The "Merging Mosaic" (Semantics-Aware Token Merging)

The Problem: Pictures are made of tiny patches (tokens). Old methods use "Hard Pruning," which is like taking a mosaic picture and smashing the tiles you don't like. If you smash the tiles that form the edge of a cat's ear, the picture of the cat is ruined.
The TASM Solution: TASM uses "Soft Merging." Instead of smashing tiles, it glues similar, less important tiles together.

It treats the image like a puzzle. If two small patches are next to each other and look similar, it combines them into one slightly larger, denser patch.
Result: The picture is smaller (compressed), but the shape and structure of the image remain intact. The robot can still "see" where things are.

3. The "Smart Filing Cabinet" (Dynamic Retrieval)

The Problem: Once you compress the memory, it becomes static. If the robot gets a tricky question later, it can't go back and get the details it threw away earlier. It's like throwing away a file folder because it looked boring, only to realize you needed it for a specific question later.
The TASM Solution: TASM builds a two-tier memory system:

Core Memory: A small, fast, high-speed desk where the most important, compressed info lives.
Latent Bank: A huge, slow, deep storage room (like a basement) where the rest of the details are kept.
The Trigger: TASM has a "surprise detector." If the robot gets a question that feels different or confusing (a high "Jensen-Shannon divergence"), it knows it needs more info. It then quickly runs to the basement (Latent Bank) to fetch the specific details it needs for that question.
Result: The robot stays fast most of the time but can instantly access deep details when the situation demands it.

The Results

The paper claims that TASM is a game-changer because:

It saves massive space: It can reduce memory usage by up to 85% (fitting a library into a backpack).
It keeps the robot smart: Unlike other methods that make the robot worse at spatial tasks (like finding an object in an image) or time-based tasks (like understanding a video), TASM keeps the robot's performance nearly as good as if it had seen all the examples without compression.
It works without retraining: You don't need to teach the robot new things; you just give it this new memory management system.

In short, TASM stops the robot from "forgetting" important details by organizing its memory like a well-structured library rather than a pile of shredded paper.

Technical Summary: Task-Aware Structured Memory for Dynamic Multi-modal In-Context Learning

1. Problem Statement

Multi-modal Large Language Models (MLLMs) rely on In-Context Learning (ICL) for rapid task adaptation, particularly in the "Many-Shot" regime where performance scales with the number of demonstrations. However, this scalability is severely constrained by two factors:

Finite Context Windows & KV Cache Costs: Processing thousands of images and text tokens generates massive Key-Value (KV) caches that overwhelm GPU memory and increase latency, often leading to Out-of-Memory (OOM) errors.
Limitations of Existing Compression: Current memory compression methods (e.g., EMLoC, SnapKV) typically rely on rigid token removal (hard pruning) or sample-dependent importance estimation (e.g., answer-aware attention). These approaches introduce three critical structural failures:
- Sample-Specific Bias: Overfitting to specific demonstration answers, potentially discarding context crucial for the test query.
- Topological Destruction: Hard pruning disrupts the 2D spatial manifold of visual tokens, destroying the spatial relationships required for tasks like localization and grounding.
- Static Memory Rigidity: Once compressed, the cache cannot recover discarded information when a complex query requires fine-grained details, leading to "context saturation" and performance degradation in long sequences.

2. Methodology: TASM Framework

The authors propose TASM (Task-Aware Structured Memory), a training-free framework designed to compress memory while preserving semantic structure and enabling dynamic retrieval. The framework operates in two phases: offline compression of the support set and online inference with dynamic retrieval.

2.1 Task-Vector Guided Importance Estimation

Instead of relying on generic attention scores or answer-specific signals, TASM extracts a global Task Vector ( $\tau$ ) representing the transformation direction from questions to answers within the few-shot demonstrations.

Extraction: The vector is computed as the normalized difference between the semantic centroids of the Answer representations and the Question representations ( $\tau = \frac{\mu_A - \mu_Q}{\|\mu_A - \mu_Q\|}$ ).
Scoring: Token importance is determined by projecting key embeddings onto this task vector. This filters out noise orthogonal to the task direction.
Layer-Adaptive Gating: Recognizing that shallow layers process local features (syntax, edges) while deep layers handle global semantics, TASM employs a gating mechanism ( $\lambda(l)$ ) that interpolates between local attention scores and global task scores. This ensures local visual context is preserved in early layers while task relevance guides compression in deeper layers.

2.2 Semantics-Aware Token Merging

To avoid the topological destruction caused by hard pruning, TASM replaces token removal with Semantics-Aware Token Merging.

Graph Matching: The compression problem is formulated as a bipartite graph matching problem. Low-scoring "Source" tokens are merged into high-scoring "Sink" tokens.
Spatial Constraints: A spatial regularizer ensures that visual tokens only merge with their spatial neighbors (within a receptive field), preserving the 2D manifold structure essential for vision tasks.
Aggregation: The compressed state is a weighted centroid of the cluster, effectively acting as a dynamic pooling layer that condenses information rather than discarding it.

2.3 Information-Geometric Dynamic Retrieval

TASM structures memory into a hierarchy comprising a compact Core Memory (GPU-resident) and a larger Latent Bank (CPU-offloaded).

Trigger Mechanism: A dynamic retrieval gate is activated based on the Jensen-Shannon (JS) Divergence between the current attention distribution and the previous step's distribution. High divergence indicates the query is attending to a novel subspace, signaling a need to retrieve historical tokens.
Sparse Retrieval: When triggered, the system fetches relevant tokens from the Latent Bank to expand the effective cache for the current step, allowing the model to adapt to query complexity without maintaining the full history in VRAM.

3. Key Contributions

Task-Vector Guided Compression: A novel approach that decouples memory compression from specific samples, using a task-level direction to capture shared reasoning patterns and reduce sample bias.
Structure-Preserving Merging: The introduction of semantics-aware token merging via bipartite graph matching, which maintains the spatial topology of visual representations, addressing the failure of hard pruning in multimodal tasks.
Dynamic Hierarchical Memory: A query-adaptive architecture that balances efficiency and adaptability by dynamically retrieving from a latent bank only when distributional shifts occur, preventing static memory rigidity.

4. Experimental Results

Evaluations were conducted on nine benchmarks, including ImageNet-100, ScreenSpot, MME-RealWorld, IllusionVQA, and VideoMME, using models like Qwen2-VL, LLaVA-NeXT, and InternVL.

Performance vs. Compression: TASM achieves performance comparable to full-context Many-Shot ICL while reducing memory usage by up to 80-85%. For instance, on ImageNet-100 with 200 examples, TASM reduced the average context length from 16,264 (Full Context) to 3,485 tokens with superior accuracy.
Spatial and Temporal Robustness: Unlike EMLoC, which suffers degradation in spatial (ScreenSpot) and temporal (YouCook2) tasks due to topological destruction, TASM preserves these capabilities. On ScreenSpot, TASM outperformed EMLoC (19.5 vs. 18.3), and on YouCook2, it surpassed the full-context baseline (109.5 vs. 108.8).
Visual Needle-in-a-Haystack (V-NIAH): TASM demonstrated superior resilience in retrieving fine-grained details from long distractor sequences. At 200 images, TASM achieved 45.5% accuracy, more than doubling EMLoC's performance (18.9%) and significantly outperforming static pruning methods like FastV (7.1%).
Scalability: TASM maintains high decoding throughput (~43.8 tok/s) and low memory footprint (11.5 GB at 32k context) on A800 GPUs, whereas full-context inference hits OOM errors at 32k tokens.
Transferability: The Task Vector exhibits strong zero-shot transferability, achieving 64.2% accuracy on disjoint target classes when computed on source classes, confirming it encodes general task semantics rather than overfitting to specific samples.

5. Significance and Claims

The paper claims that TASM addresses the fundamental bottleneck of memory scalability in multi-modal ICL without requiring model retraining. By shifting from rigid, sample-dependent pruning to task-aware, structure-preserving, and dynamically accessible memory construction, TASM enables robust long-context adaptation on consumer-grade hardware.

The authors assert that their work demonstrates that structuring memory based on task relevance and semantic topology is the key to balancing efficiency with adaptability. This allows MLLMs to leverage the "Many-Shot" regime effectively, matching or exceeding full-context performance while drastically reducing computational costs and preserving the integrity of visual representations. The framework is presented as a generalizable solution applicable across various MLLM architectures.