Technical Summary: MIITA

Problem Statement

Continual learning (CL) for Small Language Models (SLMs) faces a critical bottleneck in resource-constrained deployments. While SLMs are essential for localized, low-latency, and privacy-sensitive applications, their limited parameter capacity makes them highly susceptible to catastrophic forgetting when new supervised experiences are written directly into the backbone. Existing memory-based CL approaches, primarily designed for Large Language Models (LLMs), rely on storing textual examples, summaries, or demonstrations. These methods fail in SLM contexts for two reasons:

1. Storage Inefficiency: Textual memories retain surface content alongside reusable signals, consuming excessive storage under fixed budgets.
2. Reasoning Limitations: SLMs lack the strong in-context reasoning capabilities required to effectively infer how retrieved textual memories should influence predictions.

Consequently, current methods either degrade significantly as model size decreases or fail to mitigate forgetting under tight storage constraints.

Methodology: MIITA

The authors propose MIITA (Memory-Induced Inference-Time Adaptation), a framework designed for supervised CL in SLMs under constrained persistent storage. MIITA decouples knowledge retention from the model's parameter space by storing experiences not as raw text, but as compact functional correction directions.

Core Components

1. Direction-Oriented Memory Representation:

   • Instead of storing raw examples $(x, y)$, MIITA converts each supervised experience into a correction direction $d$ in the pre-LM-head hidden space. This direction represents the gradient update required to increase the likelihood of the target output.
   • Memories are organized as prototypes $u_j = (d_j, A_j, n_j)$, where $d_j$ is the normalized correction direction, $A_j$ is a set of lightweight semantic anchors (keys) used solely for retrieval, and $n_j$ tracks the number of merged experiences.
   • This design separates functional retention (the direction) from semantic access (the anchor), allowing semantically disparate examples that induce identical corrections to be merged.

2. Budget-Aware Memory Management:

   • Prototype Assignment: New experiences are merged into existing prototypes if their correction directions are sufficiently aligned (cosine similarity > threshold $\tau_d$). Otherwise, a new prototype is created.
   • Memory Repair: When the storage budget $C_{mem}$ is exceeded, the system prunes redundant semantic anchors, merges similar prototypes, and evicts the least valuable prototype based on a value function that prioritizes directional distinctiveness and storage efficiency.

3. Inference-Time Retrieval and Adaptation:

   • Retrieval: At inference, MIITA retrieves relevant prototypes using two cues:
     • Semantic Relevance: Matching the query's semantic key against stored anchors.
     • Uncertainty-Guided Directional Proxy: Since ground truth is unavailable at inference, MIITA computes a gradient of the restricted predictive entropy to generate a "directional proxy" indicating the steepest descent for reducing uncertainty.
   • Adaptation: Retrieved directions are aggregated into a query-specific correction vector $\bar{d}_q$. This vector is applied via a gated temporary hidden-state update:
     $$p'_t = \text{softmax}(W_{LM}(h^{out}_t + \eta \alpha_t \bar{d}_q))$$
     where $\alpha_t$ is an activation gate based on the similarity between the current generation state and the query state. This adaptation is disposable; it requires no backpropagation, parameter updates, or prompt extensions.

Theoretical Insights

The paper provides a local theoretical analysis establishing that:

• Correction directions are first-order optimal for decreasing supervised loss.
• Uncertainty-guided directions identify the steepest path for reducing local predictive entropy, serving as effective label-free retrieval cues.
• Retention of old-stage knowledge is controlled by the directional coverage of the stored memory, rather than the coverage of raw examples.

Key Contributions

1. Direction-Oriented Memory: A novel representation storing supervised experiences as correction-direction prototypes with lightweight semantic anchors, overcoming the storage and reasoning costs of textual memory in SLMs.
2. Inference-Time Adaptation Mechanism: A method to apply retrieved historical directions via gated temporary hidden-state updates, enabling non-destructive reuse of past supervision without test-time backpropagation.
3. Theoretical Analysis: A local theoretical framework linking the design to first-order loss reduction, uncertainty-guided retrieval, and the role of directional coverage in retaining knowledge.
4. Empirical Validation: Comprehensive evaluation demonstrating consistent improvements in final performance and reduced forgetting across diverse supervised CL settings and model scales.

Experimental Results

Experiments were conducted on four benchmarks (Banking77, DSC, PAXQA, SuperNI) using SLMs (Qwen3-0.6B to 4B, LLaMA-3.2-3B to 8B) under a fixed 1% memory budget.

• Performance: MIITA consistently achieved the best Overall Performance (OP) and Backward Transfer (BWT) across all benchmarks and model sizes, outperforming both parameter-update baselines (e.g., LoRA) and memory-based LLM baselines (e.g., CIS, FOREVER).
• Model Scaling: Unlike LLM-centric memory methods which degrade significantly as model size decreases, MIITA maintained robust performance even on the smallest 0.6B models.
• Budget Sensitivity: MIITA demonstrated superior performance in low-budget regimes (0.1%–2%), whereas replay-based methods required significantly larger budgets to compete.
• Ablation Studies:
  • Functional vs. Semantic: Storing functional correction directions (DirectionMem) significantly outperformed storing raw examples, summaries, or semantic embeddings, confirming that functional redundancy is distinct from semantic redundancy.
  • Retrieval: Combining semantic and uncertainty-guided retrieval yielded the best results, validating the utility of the label-free directional proxy.

Significance

The paper argues that for SLMs, the key to continual learning is not merely the use of memory, but the efficiency and functionality of the stored representation. By shifting from preserving raw text to preserving compact functional correction signals, MIITA addresses the dual bottleneck of limited storage and weak in-context reasoning. This approach allows SLMs to continually adapt to evolving needs without overwriting scarce parameter space or relying on destructive fine-tuning, offering a viable path for deploying adaptable AI in resource-constrained environments.