Transformer Configuration Details for TimesFM 2.5: Complete Architecture Guide

TimesFM 2.5 (also referred to as TimesFM 2.S) implements a 20-layer stacked transformer architecture with 1280-dimensional hidden states, 16 attention heads, RMSNorm normalization throughout, rotary positional embeddings (RoPE), and bias-free Swish-activated feed-forward networks, all centralized in the TransformerConfig dataclass.

The google-research/timesfm repository contains the official implementation of TimesFM 2.5, a foundation model for time series forecasting. Understanding the transformer configuration details for TimesFM 2.5 is essential for fine-tuning, extending, or debugging the model, as every architectural hyperparameter is explicitly defined in the source code and exposed through the TimesFM_2p5_200M_Definition class.

Core Transformer Configuration Parameters

TimesFM 2.5 defines its architecture through the TransformerConfig dataclass located in src/timesfm/configs.py. The concrete instantiation for the 200M parameter model appears in TimesFM_2p5_200M_Definition within src/timesfm/timesfm_2p5/timesfm_2p5_base.py (lines 101–117).

Parameter	Value	Description
model_dims	1280	Hidden dimension size for all transformer layers.
hidden_dims	1280	Intermediate size of the feed-forward network (FFN).
num_heads	16	Number of attention heads (head dimension = 80).
num_layers	20	Total stacked transformer blocks in the model.
attention_norm	"rms"	Applies RMSNorm before and after the attention block.
feedforward_norm	"rms"	Applies RMSNorm before and after the FFN block.
qk_norm	"rms"	Applies RMSNorm to query and key vectors within attention heads.
use_bias	False	All linear projections (attention and FFN) operate without bias terms.
use_rotary_position_embeddings	True	Enables RoPE for Q/K vectors before attention computation.
ff_activation	"swish"	Uses the Swish non-linearity in the FFN.
fuse_qkv	True	Enables fused QKV projection for optimized JAX execution.

Configuration Implementation in the Codebase

Model Definition in timesfm_2p5_base.py

The TimesFM_2p5_200M_Definition class assembles the full model by stacking 20 identical transformer layers. The configuration is passed through StackedTransformersConfig:

# src/timesfm/timesfm_2p5/timesfm_2p5_base.py (lines 101-117)

stacked_transformers = StackedTransformersConfig(
    num_layers=20,
    transformer=TransformerConfig(
        model_dims=1280,
        hidden_dims=1280,
        num_heads=16,
        attention_norm="rms",
        feedforward_norm="rms",
        qk_norm="rms",
        use_bias=False,
        use_rotary_position_embeddings=True,
        ff_activation="swish",
        fuse_qkv=True,
    ),
)

Transformer Block Implementation in transformer.py

The Transformer class in src/timesfm/flax/transformer.py (lines 291–357) consumes this configuration to build each layer using Flax NNX. The constructor initializes:

• Pre-attention RMSNorm (pre_attn_ln) and post-attention RMSNorm (post_attn_ln)
• MultiHeadAttention with RoPE, per-dimension scaling, and Q/K RMSNorm
• Pre-FFN RMSNorm (pre_ff_ln) and post-FFN RMSNorm (post_ff_ln)
• Bias-free linear layers ff0 (1280→1280) and ff1 (1280→1280) with Swish activation

Step-by-Step Architecture Flow

During the forward pass, each of the 20 transformer blocks executes the following residual operations:

1. Pre-attention RMSNorm normalizes the input before feeding it into the multi-head attention mechanism.
2. Multi-Head Attention computes attention using rotary positional embeddings on Q/K, RMSNorm on Q/K (via qk_norm), and a causal mask generated by make_attn_mask.
3. Post-attention RMSNorm stabilizes the residual stream after the attention output is added.
4. Pre-FFN RMSNorm prepares the signal for the feed-forward network.
5. FFN Computation passes through ff0 → Swish activation → ff1, with no bias terms.
6. Post-FFN RMSNorm finalizes the block before the residual addition returns control to the next layer.

Practical Code Examples

Inspecting the Default Configuration

Retrieve the exact hyperparameters from the model definition:

from timesfm.timesfm_2p5.timesfm_2p5_base import TimesFM_2p5_200M_Definition

cfg = TimesFM_2p5_200M_Definition.stacked_transformers.transformer
print("Model dim :", cfg.model_dims)        # 1280

print("Heads     :", cfg.num_heads)         # 16

print("FFN hidden:", cfg.hidden_dims)       # 1280

print("Rotary PE :", cfg.use_rotary_position_embeddings)  # True

Building a Single Transformer Layer

Instantiate a standalone transformer block matching the TimesFM 2.5 specification:

import jax
import jax.numpy as jnp
from flax import nnx
from timesfm.flax.transformer import Transformer
from timesfm.configs import TransformerConfig

config = TransformerConfig(
    model_dims=1280,
    hidden_dims=1280,
    num_heads=16,
    attention_norm="rms",
    feedforward_norm="rms",
    qk_norm="rms",
    use_bias=False,
    use_rotary_position_embeddings=True,
    ff_activation="swish",
    fuse_qkv=True,
)

rngs = nnx.Rngs(0)
transformer = Transformer(config, rngs=rngs)

# Dummy input: batch=2, seq_len=64, dim=1280

x = jnp.zeros((2, 64, 1280))
patch_mask = jnp.zeros((2, 64), dtype=bool)

out, _ = transformer(x, patch_mask)
print(out.shape)   # (2, 64, 1280)

Stacking the Full 20-Layer Model

Construct the complete stacked architecture programmatically:

from timesfm.flax.transformer import Transformer
from timesfm.configs import StackedTransformersConfig, TransformerConfig

stack_cfg = StackedTransformersConfig(
    num_layers=20,
    transformer=TransformerConfig(
        model_dims=1280,
        hidden_dims=1280,
        num_heads=16,
        attention_norm="rms",
        feedforward_norm="rms",
        qk_norm="rms",
        use_bias=False,
        use_rotary_position_embeddings=True,
        ff_activation="swish",
        fuse_qkv=True,
    ),
)

rngs = nnx.Rngs(0)
layers = [Transformer(stack_cfg.transformer, rngs=rngs) 
          for _ in range(stack_cfg.num_layers)]

def forward(x, mask):
    cache = None
    for layer in layers:
        x, cache = layer(x, mask, decode_cache=cache)
    return x

Summary

• 20 transformer layers form the backbone of TimesFM 2.5, each with 1280 model dimensions and 1280 FFN hidden dimensions.
• 16 attention heads split the hidden space into 80-dimensional heads, utilizing RMSNorm for attention, feed-forward, and Q/K normalization.
• Rotary positional embeddings provide positional information, while bias-free linear layers and Swish activation optimize the FFN.
• Configuration is centralized in TransformerConfig and instantiated via TimesFM_2p5_200M_Definition in src/timesfm/timesfm_2p5/timesfm_2p5_base.py.

Frequently Asked Questions

What is the difference between TimesFM 2.5 and TimesFM 2.S?

They refer to the same model. The source code uses "2.S" (indicating version 2.5) in identifiers like TimesFM_2p5_200M_Definition, while public documentation and release notes refer to it as TimesFM 2.5.

Why does TimesFM 2.5 use RMSNorm instead of LayerNorm?

The configuration explicitly sets attention_norm, feedforward_norm, and qk_norm to "rms". RMSNorm offers similar stabilization benefits to LayerNorm with reduced computational overhead, which the Transformer class implements in src/timesfm/flax/transformer.py using RMSNorm modules before and after each sub-layer.

How many parameters does the TimesFM 2.5 transformer have?

The TimesFM_2p5_200M_Definition class implements approximately 200 million parameters, derived from the 20 stacked layers, 1280-dimensional hidden states, and corresponding feed-forward networks.

Can I modify the transformer configuration for custom training?

Yes. You can instantiate a custom TransformerConfig with alternative model_dims, num_heads, or num_layers values and pass it to the Transformer constructor or StackedTransformersConfig. Note that pre-trained checkpoints released by Google Research are only compatible with the default 1280-dimensional, 20-layer configuration.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}