A have a look at Apple’s new Transformer-powered predictive textual content mannequin

New York, NY — September 08, 2023

At WWDC earlier this 12 months, Apple introduced that upcoming variations of iOS and macOS would ship with a brand new function powered by “a Transformer language model” that can give customers “predictive textual content suggestions inline as they kind.”

Upon listening to this announcement, I used to be fairly inquisitive about how this function works.
Apple hasn’t deployed many language fashions of their very own, regardless of most of their opponents going all-in on massive language fashions during the last couple years.
I see this on account of Apple usually priding themselves on polish and perfection, whereas language fashions are pretty unpolished and imperfect.

In consequence, this can be one of many first Transformer-based fashions that Apple will ship in one in all its working methods, or a minimum of one of many first that they’ve acknowledged publicly.
This left me with some questions in regards to the function, notably:

• What underlying mannequin is powering this function?
• What’s its structure?
• What information was used to coach the mannequin?

After spending a while with these questions, I used to be capable of finding some solutions, however most of the particulars nonetheless stay unclear.
Should you’re capable of get any additional than I may, please get in contact!

How does the function work?

After putting in the macOS beta, I instantly opened the Notes app and began typing.
Regardless of attempting many alternative sentence buildings, the function usually appeared much less usually than I anticipated it to.
It principally completes particular person phrases.

The function will sometimes recommend a couple of phrase at a time, however that is usually restricted to situations the place the upcoming phrases are extraordinarily apparent, much like the autocomplete in Gmail.

Can we dig deeper?

Discovering the mannequin itself was a bit robust, however I ultimately discovered the mannequin being utilized by AppleSpell, an inner macOS software that checks for spelling and grammar errors as you kind.
With the assistance of xpcspy, I wrote a Python script that snoops on AppleSpell exercise and streams probably the most possible recommendations from the predictive textual content mannequin as you kind in any software.

Sadly, I wrote this script earlier in the summertime, on the primary macOS Sonoma beta.
In one of many subsequent betas (I’m undecided which), Apple eliminated the unused completions from the XPC messages despatched by AppleSpell.
I wasn’t capable of glean an excessive amount of in regards to the mannequin’s habits from these completions, nevertheless it was nonetheless a cool discover.

The place is the mannequin?

After some extra digging, I’m fairly positive I discovered the predictive textual content mannequin in /System/Library/LinguisticData/RequiredAssets_en.bundle/AssetData/en.lm/unilm.bundle.
The bundle incorporates a number of Espresso mannequin information which are used whereas typing (Espresso appears to be the inner identify for the a part of CoreML that runs inference on fashions).
I wasn’t finally capable of reverse-engineer the mannequin, however I’m pretty assured that is the place the predictive textual content mannequin is stored.
Right here’s why:

1. Lots of the information in unilm.bundle don’t exist on macOS Ventura (13.5), however they do exist on the macOS Sonoma beta (14.0). And the information that do exist in each variations have all been up to date in Sonoma.
2. sp.dat, one of many information in unilm.bundle, exists on Ventura, nevertheless it’s been up to date within the Sonoma beta. Within the up to date model of the file, I discovered what appears fairly clearly like a set of tokens for a tokenizer.
3. The variety of tokens in sp.dat matches the form of the output layer in each unilm_joint_cpu.espresso.form and unilm_joint_ane.espresso.form (ANE = Apple Neural Engine), two information in unilm.bundle that describe the shapes of layers in an Espresso/CoreML mannequin. That is what we might count on to see for a mannequin that’s skilled to foretell the subsequent token.

The predictive textual content mannequin’s tokenizer

I discovered a set of 15,000 tokens in unilm.bundle/sp.dat that fairly clearly appear like they kind the vocabulary set for a big language mannequin.
I wrote a script that you should use to see this vocabulary file for your self, which you’ll take a look at on GitHub.

The vocabulary begins with <pad>, <s>, </s>, and <unk> tokens, that are all pretty widespread particular tokens (roberta-base and t5-base are two fashionable language fashions):

>>> from transformers import AutoTokenizer
>>>
>>> tokenizer = AutoTokenizer.from_pretrained("roberta-base")
>>> tokenizer.convert_ids_to_tokens([0, 1, 2, 3])
['<s>', '<pad>', '</s>', '<unk>']
>>>
>>> tokenizer = AutoTokenizer.from_pretrained("t5-base")
>>> tokenizer.convert_ids_to_tokens([0, 1, 2])
['<pad>', '</s>', '<unk>']

Subsequent come the next sequences:

• 20 particular tokens, named UniLMCTRL0 by way of UniLMCTRL19
• 79 contractions (I’d, couldn’t, you’ve…)
• 1 particular _U_CAP_ token
• 20 particular tokens, named _U_PRE0_ by way of _U_PRE19_
• 60 particular tokens, named _U_NT00_ by way of _U_NT59_
• 100 emojis

After which comes a extra normal-looking checklist of 14,716 tokens, most of that are adopted by the particular character ▁ (U+9601), which is often utilized in byte-pair encoding (BPE) tokenizers, such because the GPT-2 tokenizer, to indicate an area.

I’ve to say that this vocabulary file strikes me as fairly distinctive, nevertheless it’s positively not out of the query for a language mannequin deployed on this setting.
I’ve personally by no means seen emojis featured so prominently in a language mannequin’s tokenizer, however existing research has proven that domain-specific fashions and tokenizers can drastically enhance downstream mannequin efficiency.
So it is sensible {that a} mannequin skilled to be used in issues like textual content messages, wherein emojis and contractions might be used rather a lot, would prioritize them.

Mannequin structure

Primarily based on the contents of the unilm_joint_cpu mannequin from earlier, we are able to make some assumptions in regards to the predictive textual content community.
Regardless of sharing the identify of Microsoft’s UniLM from 2019, it appears extra to me like a mannequin based mostly on GPT-2.

GPT-2 has 4 most important elements: token embeddings, positional encodings, a sequence of 12-48 decoder blocks, and an output layer.
The community described by unilm_joint_cpu seems to be the identical, besides with solely 6 decoder blocks.
A lot of the layers inside every decoder block have names like gpt2_transformer_layer_3d, which might additionally appear to recommend it’s based mostly on a GPT-2 structure.

From my calculations based mostly on sizes of every layer, Apple’s predictive textual content mannequin seems to have about 34 million parameters, and it has a hidden measurement of 512 models.
This makes it a lot smaller than even the smallest model of GPT-2.

For the restricted scope of the predictive textual content function, this is sensible to me.
Apple needs a mannequin that may run in a short time and really often, with out draining a lot of your machine’s battery.
After I was testing the predictive textual content function, recommendations appeared virtually immediately as I typed, making for an excellent consumer expertise.
Whereas the mannequin’s restricted measurement means it wouldn’t be excellent at writing full sentences or paragraphs, when it displays very excessive confidence within the subsequent phrase or two, they’re more likely to be ok to recommend to the consumer.

See Also

Ant habits evokes autonomous materials meeting analysis

Nonetheless, with my script that snoops on exercise from AppleSpell, we are able to get the mannequin to write down full sentences anyway.
If I kind “At this time” as the primary phrase of my sentence and take the mannequin’s high suggestion every time, right here’s what I get (video):

At this time is the day of the day and the day of the week goes to be factor I’ve to do is get a brand new one for the subsequent couple weeks and I believe I’ve plenty of…

Not very inspiring.
We will evaluate this with the output from the smallest GPT-2 mannequin:

At this time, the White Home is constant its efforts in opposition to Iran to assist the brand new President, however it’ll additionally attempt to construct new alliances with Iran to make extra…

Or the most important GPT-2 mannequin:

At this time, the U.S. Division of Justice has filed a lawsuit in opposition to the town of Chicago, the Chicago Police Division, and the town’s Impartial Police Assessment Authority, alleging that the police division and the Impartial Police Assessment Authority engaged in a sample or apply…

Fairly cool seeing the consequences of all these additional parameters!
It’ll be attention-grabbing to see how this function grows and evolves sooner or later, and whether or not Apple decides to maintain its scope pretty slim or sometime broaden its skills.

Should you’re interested by attempting any of this out for your self, all of my code is on GitHub.