Skip to main content

TinyModel

TinyModel is a 4 layer, 44M parameter model trained on TinyStories V2 for mechanistic interpretability. It uses ReLU activations and no layernorms. It comes with trained SAEs and transcoders.

It can be installed with pip install tinymodel

from tiny_model import TinyModel, tokenizer

lm = TinyModel()

# for inference
tok_ids, attn_mask = tokenizer(['Once upon a time', 'In the forest'])
logprobs = lm(tok_ids)

# Get SAE/transcoder acts
# See 'SAEs/Transcoders' section for more information.
feature_acts = lm['M1N123'](tok_ids)
all_feat_acts = lm['M2'](tok_ids)

# Generation
lm.generate('Once upon a time, Ada was happily walking through a magical forest with')

# To decode tok_ids you can use
tokenizer.decode(tok_ids)

It was trained for 3 epochs on a preprocessed version of TinyStoriesV2. Pre-tokenized dataset here. I recommend using this dataset for getting SAE/transcoder activations.

SAEs/transcoders

Some sparse SAEs/transcoders are provided along with the model.

For example, acts = lm['M2N100'](tok_ids)

To get sparse acts, choose which part of the transformer block you want to look at (currently sparse MLP/transcoder and SAEs on attention out are available, under the tags 'M' and 'A' respectively). Residual stream and MLP out SAEs exist, they just haven't been added yet, bug me on e.g. Twitter if you want this to happen fast.

Then, add the layer. A sparse MLP at layer 2 would be 'M2'. Finally, optionally add a particular neuron. For example 'M0N10000'.

Tokenization

Tokenization is done as follows:

  • the top-10K most frequent tokens using the GPT-NeoX tokenizer are selected and sorted by frequency.
  • To tokenize a document, first tokenize with the GPT-NeoX tokenizer. Then replace tokens not in the top 10K tokens with a special [UNK] token id. All token ids are then mapped to be between 1 and 10K, roughly sorted from most frequent to least.
  • Finally, prepend the document with a [BEGIN] token id.

Release files for tinymodel 0.1.1.post11

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for tinymodel 0.1.1.post11
File Size Uploaded
tinymodel-0.1.1.post11.tar.gz 83.0 kB Details

Built distribution (wheel)

Table of built distributions (wheels) for tinymodel 0.1.1.post11
File Interpreter ABI Platform
tinymodel-0.1.1.post11-py3-none-any.whl Python 3 none any Details

Total release size: 159.0 kB

Release files / tinymodel-0.1.1.post11.tar.gz

Download URL tinymodel-0.1.1.post11.tar.gz
Size 83.0 kB
Tags Source
SHA-256 checksum
How to use checksums
04d5dd3af0ce4b47b68abe10ab1b7f8a0fe6849378e4ed85984eb9ea3a43a1d9
BLAKE2b-256 checksum
How to use checksums
a3d1c5415d8c8a396899157c43aa4cc87416960c8680fd79612f2f0dbd4696b5
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via poetry/1.8.3 CPython/3.12.2 Darwin/23.4.0

Release files / tinymodel-0.1.1.post11-py3-none-any.whl

Download URL tinymodel-0.1.1.post11-py3-none-any.whl
Size 76.0 kB
Tags Python 3
SHA-256 checksum
How to use checksums
ec413e8de8a1c849b890f5095a9a383253d3a919a12cb1876a4e19489a7fc11d
BLAKE2b-256 checksum
How to use checksums
5ce803b15f51f32bb9a9d459c5c485b513019c565586dbed69c9a64591cd90eb
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via poetry/1.8.3 CPython/3.12.2 Darwin/23.4.0
Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page