Zoology: Measuring and Improving Recall in Efficient Language Models

[Github]

[arXiv](2023/12/08 version v1)

 

 

Abstract

Gated Convolution-Attention Hybrid model을 통해 multi-query associative recall (MQAR) 작업에서의 성능 향상

 

 

 

Background and Preliminaries

Sequence mixers

 

Attention

 

Gated-Convolutions (= State Space Model의 일종, ≠ Mamba)

Gated-convolution의 예: Hippo, Hungry Hungry Hippos (H3)

 

 

 

Identifying the associative recall problem

Fine-grained analysis of downstream quality

Sequence mixers에 대한 언어 모델링에서의 perplex 격차. Attention이 압도적으로 좋다. (아래 3개가 gated conv)

 

Associative Recall Perplexity

AR (Associative Recall): 연관된 단어를 기억하는 것. e.g. 하쿠나 마타타

 

The pile dataset을 둘로 분리:

1. AR Hits (6.4% of tokens) : In-context에 나타났지만 훈련 중에는 1250개 이하였던 2-gram의 뒤쪽 토큰(Hakuna Matata에서 Matata).
2. Other tokens (93.6% of tokens) : 1250개 넘게 등장한 토큰

 

AR Hits는 전체 토큰의 6.4%지만 attention, gated conv 간의 품질 격차의 82%를 차지했다.

 

Formalizing the problem: Multi-Query Associative Recall

하지만 실제로 gated conv는 단일 AR에 강한데,

 

우리가 주목해야 할 것은 Multi-Query Associative Recall (MQAR) 이다.

 

 

 

Explaining the associative recall problem

BaseConv: a minimal gated convolution operator

간단하게 구현할 수 있는 layer l에 대한 최소한의 gated conv인 BaseConv를 정의한다.

입력 u, learnable filter h, linear projection W, bias b

 

Empirical analysis of gated convolution capacity and associative recall

 

Attention은 일정한 모델 차원으로 모든 시퀀스 길이의 MQAR을 완벽하게 해결한 반면 gated conv는 시퀀스가 확장됨에 따라 동일하게 선형적인 모델 확장이 필요했다. (위쪽)

 

Input-dependent filter를 사용하면 gated conv의 성능이 향상된다. 

간단한 방법 중 하나는 약간의 attention을 gated conv에 도입하는 것이다. (아래쪽)

 

 

 

Closing the Associative Recall Gap

Sparse BaseConv-Attention Hybrids

BaseConv layers와 attention layers가 있는 hybrid model을 평가한다.

일부 attention layers에 대해 선택 함수 f를 적용한다.

f의 4가지 options:

• Full attention: f(u) = 1
• Random selection
• Programmatic selection: 토큰 x가 시퀀스 이전에 등장했다면 1, 아니면 0
• Learned selection:

 

Downstream Evaluations

Attention-only model을 능가하였다.