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

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

Built distribution (wheel)

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

Total release size: 155.9 kB

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

Download URL tinymodel-0.1.1.post12.tar.gz
Size 78.5 kB
Tags Source
SHA-256 checksum
How to use checksums
e9e99a5231c80df10063efc5e588f983df8e09c3206f4a2fb9a11bbe108e6f9a
BLAKE2b-256 checksum
How to use checksums
a6f72945562f09dea1e8972353e7075527d1b0a0466c3c67a3aa499092abf7a2
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via poetry/1.8.3 CPython/3.10.12 Linux/6.5.0-41-generic

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

Download URL tinymodel-0.1.1.post12-py3-none-any.whl
Size 77.4 kB
Tags Python 3
SHA-256 checksum
How to use checksums
1a12d52dc79e60f9f19bb9add7891beb7045d73cbfadf0398c4a123114d91a39
BLAKE2b-256 checksum
How to use checksums
1d370a0300742b75770d45c28b0095df31bb4e7ae902c84c7499b6873447fd39
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via poetry/1.8.3 CPython/3.10.12 Linux/6.5.0-41-generic
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