State Soup: In-Context Skill Learning, Retrieval and Mixing (2406.08423v1)

Abstract: A new breed of gated-linear recurrent neural networks has reached state-of-the-art performance on a range of sequence modeling problems. Such models naturally handle long sequences efficiently, as the cost of processing a new input is independent of sequence length. Here, we explore another advantage of these stateful sequence models, inspired by the success of model merging through parameter interpolation. Building on parallels between fine-tuning and in-context learning, we investigate whether we can treat internal states as task vectors that can be stored, retrieved, and then linearly combined, exploiting the linearity of recurrence. We study this form of fast model merging on Mamba-2.8b, a pretrained recurrent model, and present preliminary evidence that simple linear state interpolation methods suffice to improve next-token perplexity as well as downstream in-context learning task performance.

• The paper introduces the state soup concept to interpolate internal RNN states, significantly enhancing next-token perplexity and in-context learning performance.
• The paper demonstrates a novel skill library creation using 1-shot in-context examples to efficiently generate and retrieve task-specific state vectors.
• The experiments reveal that mixing RNN states via linear averaging and A-decay techniques boosts few-shot learning, outperforming traditional sequence methods.

State Soup: In-Context Skill Learning, Retrieval, and Mixing

The paper "State Soup: In-Context Skill Learning, Retrieval and Mixing" explores a sophisticated methodology for leveraging recurrent neural networks (RNNs) to enhance efficiency in sequence modeling tasks. It introduces the concept of "state soups," a novel approach to model merging by interpolating internal states of a recurrent network, as opposed to traditional parameter interpolation. This methodology is particularly applied to a modern RNN model, Mamba-2.8B, which is pretrained for handling sequences efficiently.

RNNs, particularly those based on gated-linear architectures, have shown promise due to their ability to efficiently process long sequences—a superior trait over Transformers, whose computational costs scale quadratically with sequence length. The paper leverages this advantage by focusing on the RNNs' ability to retain and mix internal states, effectively treating them as task vectors.

Key Contributions and Methodology

1. State Soup Concept: By drawing from model soups, the paper applies a similar idea to state spaces within RNNs. The authors demonstrate that linear interpolation of states in gated-linear layers can optimize next-token perplexity and in-context learning task performance. This method not only facilitates parallel information processing but also supports efficient retrieval and adjustment of states.
2. In-Context Learning Skills: The research presents a method for creating a skill library using in-context learning tasks. These tasks are structured sequences of question-answer pairs processed by the Mamba model to output corresponding state vectors. The paper emphasizes the ability to retrieve and utilize these states effectively. Notably, the system can identify task vectors accurately with minimal input—successful cases were reported even with 1-shot examples.
3. State Mixing and Retrieval: The methodology includes mixing retrieved states to enhance learning outcomes. They explore combining states through simple linear averaging and more sophisticated techniques like A-decay mixing, which accounts for the sequential nature of data. Experimentation shows improvements in downstream few-shot learning tasks, with state mixing proving particularly effective in aligning with or outperforming traditional single-sequence processing for multi-shot examples.
4. Experimental Validation: The paper explores the retrieval and mixing approaches through a series of experiments, utilizing T-SNE for visualization and clustering analysis, demonstrating clusters aligning with specific tasks. Quantitative metrics show that state mixing enhances model performance, particularly when tasks are combined, suggesting potential for a task arithmetic model akin to arithmetic in hypervectors.

Implications and Future Directions

This paper signifies a meaningful step towards more flexible, efficient, and intuitive use of RNNs in real-world applications. By enabling the retrieval and mixing of pre-learned states, it paves the way for reuse and repurposing of these states in new contexts, offering the potential for rapid task adaptation without additional fine-tuning. The authors hint at the expanding capabilities of RNNs in handling long, many-shot tasks, forecasting a potential shift where these networks may offer advantages over Transformer-based models due to their linear cost of reusing pre-processed states.

Looking forward, the application of state soup methodology to a broader range of tasks and its ability to facilitate task arithmetic alludes to promising research avenues. Future work could extend these approaches to more dynamic and complex environments, thoroughly quantifying the advantages of this model in diverse sequence modeling tasks. Overall, this research contributes to the evolving discourse on efficient, stateful neural networks, offering a robust framework for exploitation in AI system development.