AIM: Lightweight Multi-Modal ICL

• The paper introduces AIM, a framework that transforms multi-modal demonstrations into compact virtual tokens to overcome single-image input limitations.
• AIM leverages a frozen backbone MLLM and a trainable projection layer to fuse visual information into the textual latent space, simplifying prompt complexity.
• Empirical results show significant improvements in captioning and VQA tasks while reducing memory footprint and computation by discarding excess visual tokens.

AIM is a general and light-weighted framework for multi-modal in-context learning (ICL) in multi-modal LLMs (MLLMs). It addresses two specific obstacles: most primary MLLMs are trained on single-image datasets and therefore cannot natively ingest a sequence of multi-modal demonstrations, and naive concatenation of multiple demonstrations introduces hundreds to thousands of visual tokens, causing severe memory pressure and degraded generation quality. AIM resolves these issues by aggregating image information of multimodal demonstrations into the latent space of the corresponding linguistic part, then replacing each demonstration with fused virtual tokens that are acceptable to the inner LLM, while keeping the backbone MLLM frozen (Gao et al., 2024).

1. Problem formulation and motivation

Multi-modal ICL aims to condition an MLLM on several demonstration pairs before answering a current query. In the setting discussed by AIM, the demonstrations have the form <image, prompt, label> or, more abstractly, <X_i, Y_i>, where $X_i$ is an image and $Y_i$ is the associated linguistic component. The paper identifies two concrete bottlenecks. First, most MLLMs are pretrained to accept only one image–text pair, so they cannot natively process a sequence of demonstrations. Second, when multiple image demonstrations are concatenated, the number of visual tokens grows rapidly, which stresses hardware and lengthens the context window (Gao et al., 2024).

The central empirical observation behind AIM is that the inner LLM tends to focus more on the linguistic modality within multi-modal demonstrations when generating responses. This suggests that a useful multi-modal demonstration need not be retained in its original image-heavy form, provided that the linguistic tokens can be enriched with the relevant visual information beforehand. AIM exploits exactly this asymmetry: it first lets the frozen backbone MLLM absorb image information into textual hidden states, then discards the demonstration images and keeps only a compact latent representation aligned with text.

A common misconception is that multi-modal few-shot prompting necessarily requires explicit inclusion of all demonstration images in the final prompt. AIM rejects that assumption. In its formulation, the model sees only one image at inference time—the query image—plus a set of fused tokens derived from prior image-text demonstrations. The demonstration images themselves are no longer present in the final prompt.

2. Backbone structure and fused-token construction

AIM is built around a frozen backbone MLLM such as QWen-VL or LLaVA-Next, together with its inner LLM such as QWen-7B or Vicuna-7B. The inner LLM serves a dual role: it participates in aggregating demonstration images into the linguistic space, and it remains the autoregressive generator used for answering the query (Gao et al., 2024).

For the $i$-th demonstration, the image $X_i$ is passed through the visual encoder to obtain visual tokens

$v_i \in \mathbb{R}^{m \times d},$

while the instruction plus reference label $Y_i$ is mapped to token embeddings

$x_i \in \mathbb{R}^{l \times d}.$

These are concatenated and processed by the first $K$ layers of the frozen LLM:

$$H_i = \mathrm{Transformer}_{1\ldots K}([v_i; x_i]).$$

AIM then retains only the hidden states at the textual positions,

$$f_i \coloneqq H_i[\text{text positions}] \in \mathbb{R}^{l \times d},$$

on the premise that these vectors already fuse the information of the image-text pair.

The next step is a trainable projection from these fused text hidden states to virtual tokens. With a linear adapter,

$Y_i$0

each demonstration is replaced by a sequence of $Y_i$1-dimensional virtual tokens

$Y_i$2

These tokens are acceptable to the inner LLM and function as compact surrogates for the original multi-modal demonstrations.

The architectural significance is twofold. First, AIM is “general” in the precise sense that it can be attached to any MLLM that was originally trained for single-image input, because the final prompt no longer contains multiple demonstration images. Second, it is parameter-efficient because the de facto MLLM remains frozen and only the projection layer is trained.

3. In-context learning workflow

At inference time, AIM computes the fused virtual tokens for each of the $Y_i$3 demonstrations independently, producing $Y_i$4. These are concatenated into

$Y_i$5

The frozen LLM is then fed the query image’s visual tokens, the query instruction prompt, and the concatenated fused tokens, performing ICL “as if” the prompt contained only one image (Gao et al., 2024).

This workflow is structurally important. Demonstration images are removed from the final context, but their information is preserved indirectly through text-aligned latent vectors. The inner LLM attends over the fused tokens exactly like regular text tokens, which dramatically shortens the context window. In effect, a multi-modal demonstration is nearly reduced to a pure textual demonstration, because the fused tokens stem from the textual component of the image-text pair after visual context has already been absorbed.

A plausible implication is that AIM changes the operational boundary between multi-modal prompting and text prompting. Rather than treating image demonstrations as first-class prompt objects throughout the full forward pass, it converts them into a latent textualized memory before query-time reasoning. That design choice is the mechanism by which AIM simultaneously addresses input-format incompatibility and token inefficiency.

4. Training data and optimization objective

AIM trains only the projection layer, approximately 17 M parameters, on a 56k-instance subset of the MMC4 web corpus. Each instance contains $Y_i$6 images $Y_i$7 and their matched CLIP-assigned texts $Y_i$8. For each instance, the method randomly samples $Y_i$9 image-text pairs, aggregates each $i$0 into its corresponding $i$1 to produce $i$2, and uses the remaining texts

$i$3

as the target sequence (Gao et al., 2024).

The training objective is a language-modeling loss on the remaining text:

$i$4

This is noteworthy because the training corpus has nothing to do with downstream test tasks. AIM therefore does not rely on task-specific demonstration tuning; instead, it learns a generic mechanism for converting image-text demonstrations into virtual tokens that the frozen inner LLM can exploit.

The training design also clarifies what AIM is not. It is not a full end-to-end re-training of the MLLM, and it is not a downstream adaptation method that updates billions of parameters. Its scope is narrower and more specific: learning a compact latent interface between demonstration images and the textual reasoning core.

5. Empirical performance and efficiency

The reported results focus on 16-shot ICL. With a QWen-VL 7B backbone, AIM improves Flickr30k captioning from 63.0 to 82.3 CIDEr, OKVQA from 53.3 to 56.0 VQA-acc, and VizWiz from 28.7 to 36.1 VQA-acc. On Hateful Memes in the close-ended setting, the score changes from 59.5 to 57.1 ROC-AUC. With a LLaVA-Next 7B backbone, the model without AIM fails beyond 1-shot on Flickr30k with less than 1 CIDEr, whereas AIM reaches 53.4 CIDEr at 16-shot (Gao et al., 2024).

The efficiency gains are tied directly to visual-token removal. In QWen-VL, each demonstration contributes 256 visual tokens before AIM, but after aggregation these are dropped and only $i$5 text tokens remain, giving a ratio of approximately 8%. In LLaVA-Next, 576 visual tokens per demonstration are reduced to a ratio of approximately 4%. Since memory cost scales as $i$6, discarding demonstration visual tokens keeps GPU footprint nearly constant as the number of shots grows.

Inference throughput follows the same pattern. In the 1–4-shot regime, the backbone alone is slightly faster. Beyond 8-shot, AIM becomes more efficient because the context is much shorter. This point matters because the method is not merely reducing memory usage; it changes the scaling behavior of few-shot prompting with respect to shot count.

The performance profile also sharpens a broader methodological point. AIM is not simply a compression layer inserted for convenience. Its latent reduction can improve task performance at higher shot counts because it removes the visual-token burden that otherwise degrades multi-modal ICL quality.

6. Ablations, mechanistic findings, and limitations

Ablation studies examine how many frozen LLM layers should be used to fuse image information into the text states. QWen-VL performs best when aggregation uses its first 16 layers, corresponding to half of the LLM, whereas LLaVA-Next prefers using all 32 layers to fully absorb visual cues into text (Gao et al., 2024). This indicates that the optimal depth of latent fusion is backbone-dependent rather than universal.

The paper also reports that attention maps show the frozen LLM attending predominantly to text in demonstrations. Moreover, removing visual tokens entirely in AIM often outperforms “w/o visual” backbones. The stated interpretation is that text carries most of the signal in the demonstration setting. This should not be overstated into a universal claim about all multi-modal reasoning. More precisely, within the specific AIM workflow for demonstration-conditioned generation, the textual component appears to be the dominant carrier once it has been infused with visual context.

The framework’s benefits are stated explicitly: it is general, parameter-efficient, and memory-efficient, dropping more than 90% of demonstration tokens. Its limitations are equally explicit. If labels are extremely long texts, token savings shrink because visual tokens are removed but text length remains. In addition, caching fused tokens for many demonstrations can demand storage. These limitations follow directly from the design principle of shifting demonstration information from the visual channel into a textual-length latent sequence.

Several extensions are proposed. These include quantizing or pruning the fused token bank to reduce storage, exploring nonlinear fusion such as cross-modal attention or an MLP after projection, jointly learning demonstration retrieval and aggregation end to end, and extending the same latent-space aggregation principle to video or audio modalities. This suggests that AIM is best understood not only as a specific method for image-conditioned few-shot prompting, but also as a general recipe for turning high-cost non-text demonstrations into compact LLM-native latent prompts.