FactorEngram: Factorized N-gram Memory with Basis-Level Gating for Language Models

FactorEngram. Factorized N-gram Memory with Basis-Level Gating for Language Models. By Bowen Yang and three others, from Baidu and Nanyang Technological University. Published on the twenty-eighth of September. No code repository is linked to it.

Lookup-based memory has been a promising way to scale the parameters of large language models. Instead of rebuilding the representation of a local token pattern, such as an n-gram, through layer after layer of computation, the model simply retrieves a learned representation for it. Existing designs such as Engram, though, treat each retrieved embedding as one monolithic unit. Each embedding lives in its own hashed slot, and a single scalar gate modulates it.

Two limitations follow. A polysemous pattern, one whose meaning depends on context, cannot selectively read out the parts of its memory that are relevant to the context at hand. It gets all of it, scaled up or down as a whole. And parameters are shared between patterns only when hashes happen to collide, which has largely nothing to do with meaning.

FactorEngram replaces the monolithic slot with a factorized memory. There is a dictionary of basis vectors, shared across all patterns. What a pattern retrieves is not a whole embedding but a set of coefficients over that dictionary, regularized to be sparse. Related patterns can therefore reuse common components, so sharing follows the structure of the dictionary rather than the accident of a hash.

The same dictionary is used for gating. The backbone's hidden state is scored against each basis vector, and that score gates the corresponding coefficient before the memory is reconstructed. The context therefore modulates each memory component individually, instead of turning the whole memory up or down with one number. That is exactly what the polysemous case was missing. FactorEngram also covers both individual tokens and multi-token n-grams, and the authors study systematically where in the network the memory branch should be inserted.

On Transformer backbones of three hundred and forty million parameters and of one billion parameters, FactorEngram improves language modeling and downstream task performance. Those are two different kinds of gain. Language modeling measures how well the model predicts text directly, while the downstream tasks check whether that improvement carries over into the work the model is actually used for. Ablation studies then take the design apart, and confirm that each of its components makes a contribution, so the gain is not carried by one piece while the others go along for the ride. On placement, the study identifies insertion before the attention sublayer, in the middle layers, as an effective configuration.

One dictionary, used twice. It is where patterns store what they share, and it is the axis along which context chooses what to read. Both of the old design's weaknesses, the all-or-nothing gate and the meaningless sharing, are addressed by that single change in how memory is laid out.