Text Generation Via Infini-gram


This blog post is largely inspired by one of Nathan Barry's.

Technically for n-gram models there isn't really any "training" or "learning" in the sense of training a neural network, but in this blog post I've decided to use these terms a bit more liberally.

I have an incentive to seek out alternative architecture beyond transformer-based autoregressive models because I obviously cannot compete against frontier labs with whatever resources I currently have (my last model took me 3 months and I really don't want to do it again if it's just the same thing but bigger). A few years ago when LLMs still weren't that good (unlike the monsters people are use now) some people were saying that you can do this stuff with conventional compression algorithms (since, as it turns out, language modelling is actually compression…); at the very least text classification definitely can be done without an LLM (whether it'll be accurate or not is a completely different story).

Before starting this experiment infini-gram seems to be one of the more attractive options to me purely by its name. Unbounded! Infinite! What, it seems to have the wonderful propertly of being easy to combine any number of them with near-zero cost? Brilliant! I can finally do continuous training to my heart's content without ever being punished!

Turns out it really wasn't that great at all…

Old-school n-gram

N-gram model is a surprisingly simple concept. An n-gram refers to a sequence of n words/tokens; an n-gram model generates the nth token by using the probability (approximated by frequency, which is in turn determined by the "training data") of the said token occuring with the previous n-1 tokens presents. For example, assume a 3-gram (trigram) model that uses lower-case letters as tokens and the existing sequence "a b c a b". To predict what comes next, one'd count all the occurences of "a b" and their corresponding next token in the training text (which in this case would be \(P(a|ab)\), \(P(b|ab)\), \(P(c|ab)\), \(\cdots\), \(P(z|ab)\)) and choose one of the possibilities according to the approximated probability. Naturally, an 1-gram (unigram) model use no "context" at all and only counts frequencies of every single token. Traditionally (from the literature that I found, at least) people only went as far as 5-gram, because things get prohibitively expensive very quickly; theoretically a vocabulary size of 128 and a sequence of 5 requires \(128^{5}\) probability numbers, which in half-precision (16-bit) is 64 GBytes (not to mention that you'd have to store the sequences themselves which would be another 32 GBytes). Of course one would expect that the matrix would be (relatively) sparse and you won't need to store each and every single n-gram's (estimated) probability, but the maximum \(N\) one can practically have is quite low (compared with NanoGPT, which could easily have a vocabulary size of more than 256 and an \(N\) of more than 128).

I did a bit of manual coding and implemented a simple version of n-gram models, using a subset of TinyStories as my training data. One can see from the following example that n-gram with byte-level tokenization can hardly model words (didn't happen until N >= 3), let alone sentences:

(byte-level, 1-gram) n eh dl dgan ?aoacr oei ew.sn h n aad?eeref , tdniam l t i a slols,eatusro g lueepdwo nn tde ho akin tbysaldl mohor,dfsn htihnnetwdode ge hoftisstmym rebiki h f d nlftsadoepw.eweedr b g ,dd" o.yp.t e .tsa nouatheo .i a wiloiedtnhel .e ey a ht asi.ha euh , eepeat.i h tli riuhs w enht w htd r

(byte-level, 2-gram) d, he e hilim anoo t hehitlooved itril it d pliggga wa t om be der. ithedgred n. w, ay omied y, d bled. tothearorkedan. frn falo omaysany ca tho e th ey hay? bucrittoined t ferlse tay ankemam lacl t thatoner s so s tobotoy ton wamy. f t d, asund walour w. "lieve aso try fove ba. botoowhegrees. youn

(byte-level, 3-gram w/ syncronous backoff) up ame but thappy the wittes, tastat ster mom witen. on and eards the fork, ball. ins min, a pas a name play. he spictim wer to pings cark flealkiderefuld dayin and them wastime, a puted. big up, iful the kit the rand. it theagird ch." hing big stend the thimed a le wayoutsing eved it tou ank. it k

(byte-level, 4-gram w/ syncronous backoff) p brave up and room," asked. ins a look hough use the boat toget was very knew from the will. help. sue's gree the to grass. on toy cat he lookiest and spot mouted away for pot bange. "how to set, a stoppends. the gard. i per and liked times a happy to key frossed it her sam. the six, liked son her

Syncronous backoff is a technique used here since the training data often won't cover all n-gram (some n-grams are just not valid in the language being modelled) and some n-gram would have a frequency (and thus estimated probability) of zero; in the case of a certain n-gram not having a non-zero estimated probability, we consider (n-1)-gram, (n-2)-gram, etc., until we actually have one. Naturally 1-gram is always available as an option. There is also another family of technique called "smoothing" which tries to adjust the estimation itself, but those are a bit more involved so I didn't do it here for this demonstration.

Using word-level tokens skips the stage of learning words and in my experiment has resulted in the learning of very localized language structures, but there's zero consistency across sentences whatsoever:

(word-level, 1-gram) sometimes stay at looked animals once the it . and . a . they . the weak sam the so . weapon the twist named confused gold went sure and emma and her to , he run in they was nice fun back big , grandpa inside said sister , he started they the tom wanted could was cat wore to . to prepare the girl , wanted not , shell most house her play watch some and a easy she have looked wanted her fierce his not it , fake bear and ball upon need and s she . his red moment and was special glowing ow a got salad slapped saw could and don very happy bird inside " they ' careful near so was things and girl car liked was tim expensive a saw . the a the that were . tim lily bird he ben . liked girl and they the and a run time that called girl feed a and red , keep red but " - " it asked each their and time lily and had tree a bird this them to people to and got shiny . and cats tidy " t let

(word-level, 2-gram) in his bags and said , mr . she looked at the table and he asked the top shelf . he never find somewhere in an ashtray was happy song with the bird saw blood on the shelves . they liked to eat and dad work together and mia did you help you not worry , and said , kitty liked to enjoy the children were not know what ' t take a smile . " that made her carrots . tim ' s ball all the bird with my doll . they all day , that it lived in the box was special , while amy felt the bunny liked her , even though sally , there was a lot of the soup very kind to be safe . she saw many friends , honey strip . she loved to promise to write . tim did not play with cards made an idea . they both very tough , young boy named tom took turns holding the crazy bird named lily . her kindness and said , paul felt very strong . " he had fun , amy went to fix the fruit , little creek every day

(word-level, 3-gram w/ backoff) ' t wait to eat my home pretty ! we can find ! " so , the brothers encouraged their big stack of hay at the tree , but with practice , you two , three ! " lucy wore a big sigh . the dragon . the big tree and wanted to try . what is shaking ! what do you want . " his mom were out for a little girl had an idea . the dolphin heard a noise . they were both calm as they looked inside . jane wanted to join the race and have fun . but then , a little girl smiled back . ann went to work to do . susie learned that working together with their teamwork , they decided to look for the little bird flew by them . the cat had been very kind . you will be fine . " lily felt it was dark and he wanted to see what it was out on his tight hat to kitty , and tom painted the wall for everyone to enjoy the big fish and its mommy played happily ever after .

(word-level, 4-gram w/ backoff) the monster was surprised . he did not stop crying . she was not happy to see her room was very dark . they were having so much fun . but then , a big , heavy box . the tall man closed their eyes and wished for a big playground with a seesaw . amy and her mom were baking the pie together . they looked under rocks and behind bushes . then , they saw a small crab , came to help her mom with the dishes . her mom laughed and clapped his paws . max was happy , and tim could see outside . he thought it was fun to watch them . one day , tom saw a little boy named tim . he lived in a room with a big mustache . he was not in order at all . after a long time , until he discovered a big tree and picked the sweet fruits from the garden . they dug and dug , but he did not listen to sue . they went on the water . " jane said , " it ' s max ! " tom

Infini-gram: unbounded n-gram

Infini-gram is a bit more involved because it introduces a new data structure called the suffix array. A suffix array is an array that stores the starting point of a suffix within the subject text in lexicographical order. Naturally, for any text of length \(L\) there would be \(L\) suffixes, thus the array technically would be \(L\) elements long; and since it's lexicographically sorted, all occurences of a specific substring would clumped into one continuous interval within the array; this way one can calculate the number of occurences of any specific substring by doing two binary search for the starting point and the ending point of this interval, which would be of time complexity \(O(logL)\). The big problem of directly storing the frequency values is that as context length and the size of vocabulary goes up it becomes impractical to store all the frequencies (not to mention that one should also store the n-gram sequences). The introduction of suffix array allows one to store neither while accepting having to calculate estimated probabilities on the fly.

This approach isn't without problems:

  • First of all, there would be the upfront cost of building the array, which for a dataset with \(N\) documents of average size \(M\) would have a time complexity of \(O(NM(logN + logM))\) and a space complexity of \(O(NM)\).
  • Even with the suffix array present, one'd still have to go through all the occurences to aggregate all the possible continuations to a sequence at query time if without further optimizations.
  • The text that the model's trained on must be present, since the suffix array only contains indicies.

There's also the problem of sparsity: with too small of a training dataset there would be situations where a certain (n-1)-gram has only one possibility of continuing in said dataset; if this situation repeats the model would copy & reproduce the training data verbatim. A simple way to mitigate this problem is to apply the syncronous backoff idea again: when a certain n-gram only has one particular continuation with respect to the training data, the program looks for new options with (n-1)-gram, (n-2)-gram, etc., until there's more than 1 options (naturally there'd always be more than one 1-gram options). For example, to prevent the verbatim copying of "Tom put the cup on the ground" starting from "Tom put the cup", one'd not only look for \(P(*|Tom,put,the,cup)\), but also \(P(*|put,the,cup)\), \(P(*|the,cup)\), \(P(*|cup)\) or even \(P(*|)\).

I trained a very basic version of infini-gram model using the entirety of TinyStories, and the followings are some of its outputs:

once upon a time , there were two twin cats . they lived in a small house with a nice lady . the twin cats liked to play and have fun . one day , they saw a weird big ball in the yard . the ball was very big and very red . the twin cats wanted to play with the weird big ball . they tried to push it , but it would not yield . they tried to bite it , but it was too big . the twin cats were sad , but they did not give up . they thought of a plan to play with the weird big ball . the twin cats climbed a tree and jumped on the weird big ball . the ball started to move and it went very fast . the twin cats were happy and they laughed . but the weird big ball went too fast and it went far away . the twin cats could not stop the ball and they got lost . the nice lady could not find them and they never saw their home again . <|endoftext|>

once upon a wonderful time with the ancient balloon . at the end of the day , the little boy said goodbye and watched the balloon float away until it disappeared from view . he smiled and thought it was one of the best days ever ! <|endoftext|>

once upon a wish , but he started to feel scared . suddenly , a tiny fairy appeared . the fairy said , " don ' t be scared ! tell me your wish and i ' ll make it come true ! " jack was so excited , he said , " i wish that i could have ice cream every day ! " the fairy said , " your wish is granted ! " since then , jack has enjoyed having ice cream every day . he still loves to wander around his garden looking for new adventures . <|endoftext|>

once upon a special chair just for her . tom worked very hard to make the chair . when it was done , he called sally . " sally , come see your new chair ! " he said . sally came and saw the chair . she was very happy . " thank you , tom ! " she said . and they both smiled . <|endoftext|>

once upon a plate . she said " amy , this cookie is brilliant . would you like to try it ? " amy said yes ! she took a big bite and it was so delicious . grandma said " i ' m so glad you like it ! " amy smiled and said " can we visit grandma again soon ? " everyone laughed and grandma said " yes , of course we can ! " <|endoftext|>

It might be my fault but almost none of the outputs is actually new text. I tried to tweak the parameters to the best of my ability but they always ended up being either long pieces or verbatim copies of the entries in TinyStories, and whenever there's genuine novelty it always come with the price of heavily decreased consistency.

This behaviour may be inherent to the n-gram approach:

  • There's no semantic understanding whatsoever. For example, due to "movie" and "film" being two different tokens, infini-gram would never retrieve documents containing one when the prompt only contains another (not until I somehow introduce some kind of "fuzzy equality" which I really don't have any concrete idea for), even if it really should.
  • There really isn't any mechanism that can model the causation relationship between tokens either; the only source of any consistency comes from the training data, which means one'd seriously have to choose between "consistent (but merely a copy of its training data)" and "new & novel (but makes little sense)".
    • This also makes it very sensitive to the exact location of tokens; different phrasings of the same prompt could trigger very different behaviour and generate very different results even if they have the same meaning.

Before the experiment I was actually thinking that maybe this could be used as a draft model in places that could use one (e.g. text diffusion models). Turns out it feels very much like a dead end and people should probably seek other approaches instead.

All the code can be found at https://github.com/bctnry/infini-gram-demo


Back

Last update: 2026.11.12