steklov-llama-105m

Steklov LLaMA 105M — Checkpoint Collection

Proof-of-concept checkpoints demonstrating Steklov activation sparsity on a LLaMA-style architecture.

Paper: Steklov Activations: Piecewise-Polynomial Gates with Compact Support and Tunable Sparsity Code: steklov-activations (pip install steklov-activations)

What are these models?

These are 105M-parameter LLaMA-style models (12 layers, d=768, d_ff=2048, RMSNorm, RoPE, no bias) trained on OpenWebText for 25K steps. The only difference between checkpoints is the Steklov activation scale parameter α, which controls how much of the MLP is active per token.

Checkpoints

For reference, a SiLU baseline (same architecture, no Steklov) achieves PPL 31.43 ± 0.87 with 0% activation sparsity.

Key result: The α=0.005 model has 90% of its MLP activations exactly zero on every token, yet its perplexity (30.47) is better than the dense SiLU baseline (31.43).

Downstream Benchmarks (single seed)

*SiLU baseline not included in this repo (standard LLaMA with SiLU activation).

No downstream degradation, even at 89–90% per-token activation sparsity.

How to use

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

repo = "masalskikh/steklov-llama-105m"

# Load the α=0.05 checkpoint (89% sparse, beats SiLU)
model = AutoModelForCausalLM.from_pretrained(repo, subfolder="steklov-a0.05", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(repo)

# Generate text
model.eval()
input_ids = tokenizer.encode("The future of artificial intelligence is", return_tensors="pt")
with torch.no_grad():
    for _ in range(50):
        logits = model(input_ids).logits[:, -1, :]
        next_token = torch.multinomial(torch.softmax(logits / 0.8, dim=-1), 1)
        input_ids = torch.cat([input_ids, next_token], dim=1)
print(tokenizer.decode(input_ids[0]))

# Check sparsity: count exact zeros in MLP activations
# (see steklov_llama.py get_sparsity_stats() for full profiling)

Architecture

LlamaForCausalLM(
  embed_tokens: Embedding(50257, 768)
  layers: 12 × LlamaDecoderLayer(
    self_attn: LlamaAttention(768, 12 heads)
    mlp: LlamaMLP(
      up_proj: Linear(768 → 2048)
      act_fn: SteklovSiLU(alpha=α, order=3)
      down_proj: Linear(2048 → 768)
    )
    input_layernorm: LlamaRMSNorm(768)
    post_attention_layernorm: LlamaRMSNorm(768)
  )
)

Training details

• Data: OpenWebText (2B tokens, deduplicated)
• Steps: 25,000
• Batch size: 8 × 4 grad_accum × 1024 tokens = 32K tokens/step
• Optimizer: AdamW (lr=3e-4, β₁=0.9, β₂=0.95, wd=0.1)
• Schedule: Cosine decay with 2,000 warmup steps
• Hardware: 1× RTX 5090 (multi-seed runs) / RTX 4090 (single seed)

Intended use

These checkpoints are proof-of-concept models for reproducing the paper's claims. They are not intended for production use. The 105M parameter count is too small for practical applications. Their value is in verifying:

1. Steklov activations produce exact zeros (profile the model yourself)
2. The sparsity is tunable via α
3. Quality is maintained at high sparsity
4. The 2:4 N:M compliance numbers are reproducible

Limitations

• 105M parameters (too small for practical use)
• Single-seed runs for α ≤ 0.1
• Trained for only 25K steps
• N:M sparse tensor core kernel is slower than dense at this scale
• Post-hoc activation swap does NOT work; must train from scratch

Citation

@article{masalskikh2026steklov,
  author  = {Masalskikh, A.},
  title   = {Steklov Activations: Piecewise-Polynomial Gates with Compact Support and Tunable Sparsity},
  journal = {Zenodo},
  year    = {2026},
  doi     = {10.5281/zenodo.19454642},
  url     = {https://doi.org/10.5281/zenodo.19454642}
}