about
Birdie: Advancing State Space Models with Reward-Driven Objectives and Curricula (arxiv.org)
3 points by PaulHoule on Nov 20, 2024 | hide | past | pdf | 1 comment on HN

In plain words: Instead of redesigning the network, it trains efficient models to read text both ways and pick a mix of practice tasks that pull facts from long inputs. Copying, lookup, and long-context questions improve, closing much of the gap with Transformers while staying fast and cheap.

Abstract

Efficient state space models (SSMs), such as linear recurrent neural networks and linear attention variants, offer computational advantages over Transformers but struggle with tasks requiring long-range in-context retrieval-like text copying, associative recall, and question answering over long contexts. Previous efforts to address these challenges have focused on architectural modifications, often reintroducing computational inefficiencies. In this paper, we propose a novel training procedure, Birdie, that significantly enhances the in-context retrieval capabilities of SSMs without altering their architecture. Our approach combines bidirectional input processing with dynamic mixtures of specialized pre-training objectives, optimized via reinforcement learning. We introduce a new bidirectional SSM architecture that seamlessly transitions from bidirectional context processing to causal generation. Experimental evaluations demonstrate that Birdie markedly improves performance on retrieval-intensive tasks such as multi-number phone book lookup, long paragraph question-answering, and infilling. This narrows the performance gap with Transformers, while retaining computational efficiency. Our findings highlight the importance of training procedures in leveraging the fixed-state capacity of SSMs, offering a new direction to advance their capabilities. All code and pre-trained models are available at https://www.github.com/samblouir/birdie, with support for JAX and PyTorch.

Sam Blouir, Jimmy T. H. Smith, Antonios Anastasopoulos, Amarda Shehu
arXiv:2411.01030 · cs.CL, cs.AI, cs.LG · submitted Nov 1, 2024 · updated Feb 21, 2025
abstract · pdf · html · Accepted to EMNLP 2024 (Main Conference)

add comment on HN

I don't get state space models/mamba,

Essentially, as information comes in, they have a filter to decide which information is worth remembering regardless what future tasks they will be facing. This seems to work well only when the task is fixed, such as translation.

Transformers on the other hand remember everything within a window, but when a task comes in, they can look back and fetch relevant information.

I think transformers are inherently stronger, it resembles how human think. The (fetched) memory has to be task dependent.