Asterisk-Pi-135M

Asterisk-Pi: ASPP-Attention with π-Flow Refinement

Asterisk-Pi is an enhanced version of the Asterisk model that adds π-flow (probability flow) refinement to the hybrid ASPP-Attention architecture. Building on the SmolLM2-135M base, Asterisk-Pi implements per-layer iterative refinement inspired by probability flow ODEs from diffusion models, enabling multi-step reasoning through continuous state evolution.

Model Description

• Base Model: Asterisk (SmolLM2-135M-Instruct with ASPP)
• Architecture: Hybrid ASPP-Attention + Per-Layer π-Flow (30 hybrid layers)
• Parameters: 173.7M (37.5M ASPP + 2.5M π-flow parameters)
• Training: Supervised Fine-Tuning on Mixed Benchmark Dataset
• Framework: Transformers 4.57.6, TRL 0.27.0

Key Innovation: π-Flow Refinement

π-Flow (Probability Flow) adds iterative refinement to each hybrid layer, inspired by continuous-time probability flow ODEs:

h' = h + α * v(h)  [Euler discretization]

Where:

• v(h) is the velocity field computed by a dedicated ASPP operator
• α is a learnable per-token scaling factor (adaptive gating)
• Applied after ASPP-Attention fusion in each layer

This enables 60 total refinement steps (30 layers × 2 steps each) throughout the model, allowing gradual convergence to more refined representations.

Evaluation Results

Evaluated on LM-Evaluation-Harness:

Analysis

π-Flow improvements over base Asterisk:

• ARC-Challenge (+1.54%): More challenging reasoning benefits from iterative refinement
• WinoGrande (+1.81%): Multi-step resolution helps with pronoun disambiguation

Improvements over SmolLM2-135M base:

• ARC-Challenge (+2.65%): Hybrid architecture + π-flow significantly improves complex reasoning
• ARC-Easy (+5.13%): Strong gains on elementary science questions
• WinoGrande (+2.37%): Better pronoun disambiguation through iterative refinement
• PIQA (+0.71%): Modest gains on physical commonsense

Outperforming Gemma-3-270m-it (with 96M fewer parameters):

• ARC-Challenge (+3.08%): Superior reasoning despite being 35% smaller
• ARC-Easy (+3.53%): Significant advantage on elementary science
• HellaSwag (+2.70%): Much stronger commonsense reasoning
• WinoGrande (+1.34%): Better coreference resolution
• PIQA (+0.11%): Comparable physical reasoning

Key insight: Asterisk-Pi (173.7M params) consistently outperforms the much larger Gemma-3-270m-it (270M params), demonstrating that the hybrid ASPP-Attention architecture with π-flow refinement achieves superior parameter efficiency. The structured reasoning approach enables better performance per parameter, especially on complex multi-step reasoning tasks.

Architecture

Overview

Figure: Asterisk-Pi architecture showing the hybrid ASPP-Attention structure with π-flow refinement. Each of the 30 layers contains parallel ASPP and Attention branches, gated fusion, and iterative π-flow refinement using probability flow ODE.

Input → [30 Hybrid Layers with π-Flow] → Output

Each Hybrid Layer:
1. ASPP-Attention Fusion (from base Asterisk)
2. π-Flow Refinement (NEW)
3. Feed-Forward Network

1. Hybrid ASPP-Attention Layer (Base Asterisk)

class HybridASPPAttentionLayer:
    """
    Combines ASPP operator with standard attention

    Components:
    - ASPP operator: Local structured reasoning with Union-Find graph propagation
    - Standard attention: Global context
    - Gated fusion: Dynamic balancing
    """

ASPP Operator: Union-Find Graph Propagation

The ASPP operator uses a Union-Find (Disjoint Set Union) structure for efficient graph-based message passing. Unlike traditional attention's O(n²) complexity or skip-list's O(n log n), Union-Find achieves O(n) complexity with nearly constant-time operations.

Graph Structure - Union-Find Parent Chain:

Position:  [0]  [1]  [2]  [3]  [4]  [5]  ...  [n-1]
Parent:    [0] ← 0  ← 1  ← 2  ← 3  ← 4  ...  ← n-2
           (root)

- Position 0: points to itself (root of the tree)
- Position i (i>0): points to position i-1 (parent)
- Forms a linear chain structure for sequential token relationships

This creates a directed acyclic graph (DAG) where information flows from children to parents, naturally capturing left-to-right sequential dependencies in language modeling.

Graph Propagation Aggregation:

Each ASPP evolution step performs parent-based message passing:

# Pseudocode for one ASPP propagation step
for position i in sequence:
    # 1. Find parent using Union-Find structure
    parent_idx = compute_parent_indices()[i]  # O(1) with path compression

    # 2. Gather parent features
    parent_features = hidden_states[parent_idx]

    # 3. Message aggregation: combine self + parent
    message_input = concat([hidden_states[i], parent_features])

    # 4. Update via learned transformation
    new_state = message_net(message_input)  # 2-layer MLP

    # 5. Scaled residual connection
    hidden_states[i] = hidden_states[i] + residual_scale * new_state
    hidden_states[i] = layer_norm(hidden_states[i])

Key properties of Union-Find propagation:

1. O(n) Complexity: Each position performs exactly one parent lookup and one aggregation

   • No expensive attention computation (O(n²))
   • No multi-level skip connections (O(n log n))
   • Simple indexing operation: parent_features = h[parent_indices]

2. Hierarchical Information Flow: After K steps, position i can access information from positions [i-K, i]

   • K=1: immediate parent only
   • K=2: grandparent (2 positions back)
   • K=4 (default): great-great-grandparent (4 positions back)
   • Information propagates through the chain structure

3. Learnable Aggregation: The message_net MLP learns how to combine self and parent features

   • Input: [self_features || parent_features] (2D dimensions)
   • Output: D dimensional update vector
   • Dropout regularization for robustness

4. Path Compression Potential: Can extend to dynamic parent reassignment

   • Current implementation: static parent[i] = i-1 chain
   • Future extension: learn parent assignments based on semantic similarity
   • Enables adaptive graph structure during forward pass

Union-Find vs. Other Graph Structures:

Union-Find achieves the lowest complexity while maintaining effective information propagation through iterative K-step evolution.

Theoretical Foundation - Union-Find in Graph Algorithms:

Union-Find is a classic data structure for disjoint set operations:

• Find: Determine which set an element belongs to (with path compression: O(α(n)) ≈ O(1))
• Union: Merge two sets into one
• Applications: Kruskal's MST algorithm, connected components, cycle detection

In Asterisk-Pi:

• Each token position is a node in the graph
• Parent pointers define the tree structure
• Message passing simulates "Find" operations (traversing to ancestors)
• Can extend to dynamic "Union" operations (merging related tokens)

Multi-Step Propagation:

With K=4 evolution steps, information flow becomes:

Step 1: Position i accesses parent i-1
Step 2: Position i now has information from i-2 (via i-1)
Step 3: Position i now has information from i-3 (propagated through chain)
Step 4: Position i now has information from i-4 (fully propagated)

Result: Each position has aggregated context from 4 previous positions
        through efficient O(n) operations

This multi-step propagation is crucial for:

• Local context: Recent tokens for coherence
• Gradient flow: Direct paths for backpropagation
• Efficiency: Linear cost instead of quadratic attention

Fusion mechanism:

aspp_out = ASPP(hidden_states)            # Union-Find graph propagation (O(n))
attn_out = Attention(hidden_states, mask, ...)  # Global attention (O(n²))
gate = sigmoid(linear([aspp_out || attn_out]))
fused = gate * aspp_out + (1 - gate) * attn_out

# Combines:
# - Local structured reasoning (ASPP via Union-Find)
# - Global contextual awareness (Attention)

2. π-Flow Refinement (Per-Layer)

# Added to each hybrid layer
self.pi_flow_aspp = ASPPOperator(...)        # Velocity field network
self.pi_flow_scale = Parameter(0.2)          # Learnable flow strength
self.pi_flow_gate = MLP(hidden_size -> 1)    # Token-wise adaptive gating

π-Flow forward pass:

function π_flow_refinement(hidden_states):
    for step = 1 to π_flow_steps:
        # Compute velocity field using dedicated ASPP
        v = pi_flow_aspp(hidden_states)

        # Adaptive per-token gating
        gate = sigmoid(pi_flow_gate(hidden_states))  # [B, L, 1]
        alpha = pi_flow_scale * gate

        # Euler step in probability space
        hidden_states = hidden_states + alpha * v

    return hidden_states

Key design choices:

1. Per-layer π-flow: Each of 30 layers has independent π-flow parameters
2. Learnable scale: pi_flow_scale adapts flow strength during training
3. Token-wise gating: Different tokens get different flow magnitudes
4. ASPP velocity: Reuses ASPP architecture for computing v(h)

3. Complete Layer Pseudocode

function HybridLayerWithPiFlow(hidden_states, attention_mask, ...):
    residual = hidden_states
    hidden_states = input_layernorm(hidden_states)

    # === Hybrid ASPP-Attention (Base Asterisk) ===
    aspp_output = aspp_operator(hidden_states)
    attn_output = self_attention(hidden_states, attention_mask, ...)

    # Gated fusion
    fusion_input = concat([aspp_output, attn_output])
    gate = sigmoid(linear(dropout(fusion_input)))
    fused_output = gate * aspp_output + (1 - gate) * attn_output

    # Residual connection
    hidden_states = residual + fused_output

    # === π-Flow Refinement (NEW) ===
    for step in [1..pi_flow_steps]:
        v = pi_flow_aspp(hidden_states)
        alpha = pi_flow_scale * sigmoid(pi_flow_gate(hidden_states))
        hidden_states = hidden_states + alpha * v

    # === MLP Block ===
    residual = hidden_states
    hidden_states = post_attention_layernorm(hidden_states)
    hidden_states = mlp(hidden_states)
    hidden_states = residual + hidden_states

    return hidden_states

Parameter Breakdown

π-Flow adds only 1.4% more parameters (2.5M) compared to base Asterisk (171.2M) while providing 60 total refinement steps.

Quick Start

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    "NoesisLab/Asterisk-Pi",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("NoesisLab/Asterisk-Pi")

# Generate text
messages = [{"role": "user", "content": "Explain the waterfall model in software engineering."}]
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)

outputs = model.generate(
    inputs,
    max_new_tokens=256,
    temperature=0.7,
    do_sample=True,
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Training Details

Training Dataset

Mixed benchmark dataset for testing true capabilities:

Total: ~10,148 training samples

Training Configuration

• Starting Point: Asterisk checkpoint (base ASPP-Attention model)
• Optimizer: AdamW (lr=5e-4, weight_decay=0.1)
• Batch Size: 2 per device, gradient accumulation=4 (effective batch=8)
• Epochs: 2
• Scheduler: Linear warmup (10% of steps)
• Mixed Precision: bfloat16
• Gradient Checkpointing: Enabled
• Max Grad Norm: 1.0

π-Flow Configuration

pi_flow = True
pi_flow_steps = 2           # 2 refinement steps per layer
pi_flow_scale = 1.0         # Initial flow strength
pi_flow_use_gate = True     # Token-wise adaptive gating

ASPP Configuration (Inherited from Base)

aspp_hidden_dim = 256       # Internal dimension (vs 576 model hidden_size)
aspp_num_steps = 4          # Evolution steps for ASPP
aspp_dropout = 0.2          # Regularization
hybrid_layer_indices = None # All 30 layers

Model Creation from Base Asterisk

from AsteriskForCausalLM import AsteriskForCausalLM
from safetensors.torch import load_file
import torch

# Load Asterisk config and inject π-flow parameters
from AsteriskForCausalLM import AsteriskConfig
config = AsteriskConfig.from_pretrained("path/to/Asterisk", trust_remote_code=True)

# Add π-flow configuration
config.pi_flow = True
config.pi_flow_steps = 2
config.pi_flow_scale = 1.0
config.pi_flow_use_gate = True

# Create model with π-flow
model = AsteriskForCausalLM(config)

# Load pretrained Asterisk weights (strict=False ignores new π-flow params)
state_dict = load_file("path/to/Asterisk/model.safetensors")
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)

# π-flow parameters are randomly initialized
print(f"New π-flow parameters: {len(missing_keys)}")

# Move to device
model = model.to(dtype=torch.bfloat16, device="cuda")

Theoretical Background

π-Flow: Probability Flow ODE

Inspired by diffusion model score-based formulations:

dx/dt = v(x, t)  [Continuous probability flow]

Discretized with Euler method:

x_{t+1} = x_t + Δt * v(x_t)

In Asterisk-Pi:

• x_t = hidden states at layer output
• v(x_t) = velocity field from dedicated ASPP
• Δt = learnable pi_flow_scale * gate(x_t)

Multi-Scale Refinement

• Layer-level: 30 hybrid layers with ASPP-Attention fusion
• π-Flow level: 2 steps per layer = 60 total refinement operations
• ASPP-level: 4 evolution steps within each ASPP = 240 micro-updates

This creates a hierarchical refinement cascade enabling gradual convergence to high-quality representations.

Why π-Flow Helps

1. Iterative refinement: Multiple passes allow correcting errors
2. Adaptive flow: Token-wise gating focuses computation where needed
3. Gradient flow: More direct paths for gradient propagation
4. Expressiveness: Increases model capacity with minimal parameters

Implementation Details

Return Type Handling

Critical for Transformers compatibility:

# HybridASPPAttentionLayer.forward() returns tensor only
def forward(self, hidden_states, ...) -> torch.Tensor:
    # ... ASPP + Attention + π-flow ...
    return hidden_states  # ✅ Tensor, not tuple

# This matches LlamaDecoderLayer API: -> torch.Tensor

Gradient Checkpointing Compatibility

π-Flow is fully compatible with gradient checkpointing:

• All operations are standard PyTorch ops
• No custom CUDA kernels
• Automatic differentiation through flow steps

Weight Initialization

• ASPP parameters: Transferred from base Asterisk
• π-Flow ASPP: Randomly initialized (Xavier uniform)
• π-Flow scale: Initialized to 0.2 (conservative)
• π-Flow gate: Initialized to output ~0.5 (balanced)

Files in Checkpoint

Asterisk-Pi/
├── AsteriskForCausalLM.py    # Model implementation (with π-flow)
├── config.json                # Model configuration
├── model.safetensors          # Model weights
├── tokenizer.json             # Tokenizer
├── generation_config.json     # Generation settings
└── README.md                  # This file

Differences from Base Asterisk

Known Issues & Solutions

1. Return Type Errors

Issue: AttributeError: 'tuple' object has no attribute 'dtype'

Solution: HybridASPPAttentionLayer.forward() must return torch.Tensor only, not tuple. This matches the LlamaDecoderLayer API in transformers 4.57.6.

2. π-Flow in All Layers vs Final Layer

Initial approach: π-flow only in final layer (limited expressiveness)

Current approach: π-flow in all 30 hybrid layers for maximum refinement capability.

3. Training Stability

π-Flow can cause instability with high learning rates. Use:

• Lower learning rate (5e-4 vs 2e-5 for base)
• Gradient clipping (max_norm=1.0)
• Conservative initial flow scale (0.2-1.0)

Dependencies

pip install torch>=2.0.0
pip install transformers>=4.40.0
pip install trl>=0.8.0
pip install datasets>=2.14.0
pip install accelerate>=0.25.0
pip install bitsandbytes
pip install safetensors

Citations

If you use this model, please cite:

@misc{asteriskpi2026,
  title={Asterisk-Pi: Probability Flow Refinement for Hybrid ASPP-Attention Models},
  author={NoesisLab},
  year={2026},
  publisher={Huggingface},
  url={https://huggingface.co/NoesisLab/Asterisk-Pi}
}

@misc{asterisk2026,
  title={Asterisk: Hybrid ASPP-Attention Architecture for Enhanced Language Modeling},
  author={NoesisLab},
  year={2026},
  publisher={Huggingface},
  url={https://huggingface.co/NoesisLab/Asterisk}
}

@misc{vonwerra2022trl,
  title={{TRL: Transformer Reinforcement Learning}},
  author={Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallouédec},
  year={2020},
  journal={GitHub repository},
  publisher={GitHub},
  howpublished={\url{https://github.com/huggingface/trl}}
}

@article{allal2024SmolLM2,
  title={SmolLM2 - with great data, comes great performance},
  author={Allal, Loubna Ben and Lozhkov, Anton and Penedo, Guilherme and Wolf, Thomas and von Werra, Leandro},
  year={2024}
}

Related Work

• Diffusion Models: π-flow inspired by probability flow ODEs in score-based diffusion
• Neural ODEs: Continuous-depth models with adaptive computation
• Iterative Refinement: Multi-pass decoding in sequence models

Future Directions

1. Adaptive π-flow steps: Learn number of refinement steps per layer
2. Higher-order ODE solvers: Replace Euler with RK4 or adaptive schemes
3. Stochastic π-flow: Add noise injection for exploration
4. Cross-layer π-flow: Allow information flow between distant layers

License

This model inherits the Apache 2.0 license from SmolLM2-135M-Instruct.

Framework Versions

• TRL: 0.27.0
• Transformers: 4.57.6
• PyTorch: 2.8.0+cu128
• Datasets: 4.5.0
• Tokenizers: 0.22.2

Acknowledgments

Built on top of:

• Asterisk - Base ASPP-Attention architecture
• SmolLM2-135M-Instruct - Foundation model
• TRL - Training framework

Special thanks to the diffusion model community for probability flow ODE insights.