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.post10

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.post10
File Size Uploaded
tinymodel-0.1.1.post10.tar.gz 77.0 kB Details

Built distribution (wheel)

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

Total release size: 153.0 kB

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

Download URL tinymodel-0.1.1.post10.tar.gz
Size 77.0 kB
Tags Source
SHA-256 checksum
How to use checksums
4befe2d69d8eefde1e2a4b6809ff7fe8f53c6439992ad8b22f908799d857a14f
BLAKE2b-256 checksum
How to use checksums
caadfe92eebc38281e631878cbadbdf2dd4ca66d2edc089b29208d586278d6a2
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.post10-py3-none-any.whl

Download URL tinymodel-0.1.1.post10-py3-none-any.whl
Size 76.0 kB
Tags Python 3
SHA-256 checksum
How to use checksums
671a45670090d6e0b6009b9800ebcfb484670e346f3784de0e810e07547e8b91
BLAKE2b-256 checksum
How to use checksums
8fdb27637613c5ef57175377ba6ea863c169e66ca35a66dbb490a93bd955f48c
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