README.md

metadata

license: apache-2.0
base_model: HuggingFaceTB/SmolLM2-135M
library_name: onnx
pipeline_tag: text-generation
tags:
  - transformers.js
  - onnx
  - interpretability
  - attention
  - hidden-states
  - tuned-lens

SmolLM2-135M Observable

This is a weight-only Q4 ONNX export of HuggingFaceTB/SmolLM2-135M for an interactive course laboratory. It exposes the intermediate tensors needed to inspect one forward pass in a browser.

The model is a base next-token predictor. It is not an instruction-following or chat model.

Inputs

The sequence inputs use int64 tensors with shape [batch, sequence]:

• input_ids
• attention_mask

query_index is an int64 tensor with shape [batch]. It selects the token position whose layer update will be decomposed.

The browser laboratory uses a batch size of one and short sequences.

Outputs

The graph returns 182 tensors:

• next_token_logits, with shape [batch, 49152];
• hidden_state_00 through hidden_state_30, with shape [batch, sequence, 576];
• attention_01 through attention_30, with shape [batch, 9, query sequence, key sequence];
• attention_contribution_01 through attention_contribution_30, with shape [batch, sequence, 576];
• counterfactual_hidden_state_01 through counterfactual_hidden_state_30, with shape [batch, sequence, 576];
• attention_output_01 through attention_output_30, with shape [batch, 576];
• mlp_output_01 through mlp_output_30, with shape [batch, 576].

hidden_state_00 contains the input embeddings. The following hidden states contain the raw residual stream after each transformer layer, before the final RMSNorm.

For a query token i and source token j, each contribution is:

cᵢ⟵ⱼ = Wₒ concatₕ(aᵢⱼʰ vⱼʰ)

The source contributions sum to attention_output. Adding that result to the previous residual state gives the post-attention state. Adding mlp_output gives the next residual state. The manifest records numerical checks for both identities.

Each counterfactual_hidden_state removes one source contribution before the layer MLP and evaluates that MLP again. Subtracting this state from the complete state gives the source token's effect after the nonlinearity.

Quantization

The published graph uses 4-bit weight-only quantization with asymmetric 32-value blocks. Activations and outputs remain float32. The graph uses the ONNX Runtime MatMulNBits operator.

The export script compares the ONNX result with the original PyTorch model. The included manifest records the validation prompt, tensor shapes and numerical comparison for the published file.

Intended use

This artifact supports a browser laboratory that visualizes:

• next-token probabilities;
• token boundaries;
• mean attention across heads in the selected layer;
• the most likely next token if the residual stream were read at the selected layer.

The repository includes two lens graphs. model_logit_lens_q8.onnx applies the model's final RMSNorm and original lm_head directly to an intermediate residual-stream state. model_tuned_lens_q8.onnx first applies an affine translator trained for that layer, then uses the same RMSNorm and lm_head.

The tuned lens uses 30 affine translators trained on WikiText-2. Ridge regression initializes each translator against the final residual stream. A sampled KL stage then matches the final distribution over the teacher's top 128 tokens and 128 random negatives. The training report is available in tuned-lens-training.json.

On eight held-out prompts, the tuned lens reduced mean full-vocabulary KL from 20.51 for the plain logit lens to 2.15. Mean top-10 overlap with the final distribution rose from 1.67 to 3.74 tokens. The Q8 ONNX export preserved at least 8 of the PyTorch lens's top 10 tokens on every tested layer.

Limitations

• The Q4 graph does not reproduce the float32 model exactly.
• The graph does not use a KV cache and recomputes the whole sequence when the selected token changes.
• Small base models produce weak and sometimes incoherent continuations.
• The tuned lens is a course artifact trained on WikiText-2. It is not a research-grade checkpoint or a causal explanation of the model's computation.
• Attention weights show values calculated inside the model. They do not establish a causal explanation for an output.
• This export is meant for teaching and inspection, not production inference.

Source

The export script and browser laboratory were created for the Engenharia Assistida por IA course. Hugging Face published the original SmolLM2 weights and architecture under the Apache 2.0 license.