Asterisk / README.md
OzTianlu's picture
Update README.md
5bd36fb verified
---
library_name: transformers
model_name: Asterisk
base_model: HuggingFaceTB/SmolLM2-135M-Instruct
tags:
- aspp
- hybrid-architecture
- graph-reasoning
- sft
- trl
license: apache-2.0
language:
- en
---
# Asterisk: Hybrid ASPP-Attention Architecture
**Asterisk** is a research implementation that combines the **ASPP (Adjacency-Structured Parallel Propagation)** operator with standard attention mechanisms to enhance the SmolLM2-135M model. The model implements a hybrid architecture that fuses graph-based local reasoning (ASPP) with global attention for improved expressiveness on structured reasoning tasks.
## Model Description
- **Base Model**: [SmolLM2-135M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct)
- **Architecture**: Hybrid ASPP-Attention (30 hybrid layers)
- **Parameters**: 171.2M (35M additional ASPP parameters)
- **Training**: Supervised Fine-Tuning on Capybara dataset
- **Framework**: Transformers 4.57.6, TRL 0.27.0
## Evaluation Results
Evaluated on LM-Evaluation-Harness:
| Task | Metric | Score | Stderr |
|------|--------|-------|--------|
| **HellaSwag** | acc_norm | **0.4430** | ±0.0157 |
| **ARC-Easy** | acc_norm | **0.5450** | ±0.0158 |
| **ARC-Challenge** | acc_norm | **0.2884** | ±0.0132 |
| **PIQA** | acc_norm | **0.6770** | ±0.0148 |
| **WinoGrande** | acc | **0.5210** | ±0.0158 |
### Key Innovation: The Asterisk Operator (★-operator)
The **Asterisk Operator** performs local parallel state evolution through point-wise transformations:
```
h_i^(t+1) = φ(h_i^(t)) [K-step iterative evolution]
```
This is then gated and fused with standard Llama attention outputs:
```
output = gate * ASPP(x) + (1-gate) * Attention(x)
```
## Architecture
### 1. ASPPOperator (Point-wise Parallel Propagation)
```python
class ASPPOperator:
"""
Forward pass:
1. Optional dimensionality reduction: h_t = down_proj(hidden_states)
2. K-step evolution: h_t = h_t + α * φ(h_t) [K times]
3. Layer normalization after each step
4. Optional projection back: output = up_proj(h_t)
Parameters:
- hidden_size: 576 (model dimension)
- aspp_hidden_dim: 256 (internal ASPP dimension)
- aspp_num_steps: 8 (evolution iterations)
- aspp_dropout: 0.2
"""
```
**Pseudocode:**
```
function ASPP(hidden_states):
# Optional dimensionality reduction
if use_projection:
h_t ← down_proj(hidden_states)
h_t ← dropout(h_t)
else:
h_t ← hidden_states
# Learnable number of steps
k_steps ← max(1, int(sigmoid(k_logit) * num_steps))
# K-step point-wise evolution
for t = 1 to k_steps:
# Point-wise update: φ(h_t) = MLP(h_t)
h_t_next ← update_net(h_t)
# Scaled residual connection
h_t ← h_t + residual_scale * h_t_next
h_t ← layer_norm(h_t)
# Project back to original dimension
if use_projection:
h_t ← up_proj(h_t)
h_t ← dropout(h_t)
return h_t
```
### 2. HybridASPPAttentionLayer
```python
class HybridASPPAttentionLayer(LlamaDecoderLayer):
"""
Extends LlamaDecoderLayer with parallel ASPP branch
Architecture:
1. Input LayerNorm
2. Parallel branches:
- ASPP operator for local structured reasoning
- Standard LlamaAttention for global context
3. Gated fusion: gate * ASPP + (1-gate) * Attention
4. Residual connection
5. Feed-forward MLP
"""
```
**Pseudocode:**
```
function HybridLayer(hidden_states, attention_mask, ...):
residual ← hidden_states
hidden_states ← input_layernorm(hidden_states)
# Parallel branches
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
# 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
```
### 3. AsteriskForCausalLM
```python
class AsteriskForCausalLM(LlamaForCausalLM):
"""
Main model class with custom model_type "asterisk"
Configuration:
- hybrid_layer_indices: None (all 30 layers are hybrid)
- aspp_hidden_dim: 256 (reduces overfitting)
- aspp_num_steps: 8 (learnable, actual steps ≈ 6)
- aspp_dropout: 0.2
"""
```
**Note**: These are preliminary results with sample limits. Full evaluation pending.
## Quick Start
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
"path/to/Asterisk",
trust_remote_code=True,
torch_dtype=torch.bfloat16,
device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("path/to/Asterisk")
# Generate text
messages = [{"role": "user", "content": "Explain quantum computing in simple terms."}]
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 Configuration
- **Dataset**: Capybara (conversational instruction-following)
- **Optimizer**: AdamW (lr=2e-5, weight_decay=0.01)
- **Batch Size**: 4 per device, gradient accumulation=4 (effective batch=16)
- **Epochs**: 2
- **Scheduler**: Cosine with warmup (100 steps)
- **Mixed Precision**: bfloat16
- **Gradient Checkpointing**: Enabled
### ASPP Configuration
```python
aspp_hidden_dim = 256 # Internal dimension (vs 576 model hidden_size)
aspp_num_steps = 8 # Max evolution steps (learnable)
aspp_dropout = 0.2 # Regularization
hybrid_layer_indices = None # All 30 layers
```
## Model Creation from Base
```python
from AsteriskForCausalLM import AsteriskForCausalLM
# Create Asterisk model from SmolLM2 base
model, base_model = AsteriskForCausalLM.from_pretrained_base(
"HuggingFaceTB/SmolLM2-135M-Instruct",
hybrid_layer_indices=None, # None = all layers
aspp_hidden_dim=256, # Internal ASPP dimension
aspp_num_steps=8, # K-step evolution
aspp_dropout=0.2, # Dropout rate
torch_dtype=torch.bfloat16,
device_map="auto",
)
# Base model parameters are transferred, ASPP parameters initialized randomly
model.load_state_dict(base_model.state_dict(), strict=False)
```
## Theoretical Background
### Universality (Theorem 2.1)
ASPP can simulate any Message-Passing Neural Network (MPNN) function on finite graphs in D steps, where D is the graph diameter.
### Convergence (Theorem 2.2)
Exponential convergence to fixed points with rate c=0.76 under Lipschitz continuity.
### Turing Completeness
Proven via cyclic tag system simulation - ASPP can compute any Turing-computable function given sufficient depth.
**Implementation Note**: This implementation simplifies theoretical ASPP to point-wise evolution to reduce overfitting while maintaining iterative refinement benefits.
## Files in Checkpoint
```
Asterisk/
├── AsteriskForCausalLM.py # Model implementation (required for trust_remote_code)
├── config.json # Model configuration with auto_map
├── model.safetensors # Model weights
├── tokenizer.json # Tokenizer
├── generation_config.json # Generation settings
└── README.md # This file
```
## Dependencies
```bash
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
```
## Citations
If you use this model, please cite:
```bibtex
@misc{asterisk2026,
title={Asterisk: Hybrid ASPP-Attention Architecture for Enhanced Language Modeling},
author={NoesisLab},
year={2026},
publisher={Huggingface},
url={https://huggingface.co/NoesisLab/Asterisk}
}
```
```bibtex
@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}}
}
```
```bibtex
@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}
}
```
## 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 [SmolLM2-135M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct) by HuggingFace. Training framework powered by [TRL](https://github.com/huggingface/trl).