seren-llm 0.1.1

Seren: A Self-Refining Decoder-Only LLM with GQA, RoPE, Flash Attention, and a novel two-pass self-refinement mechanism.

Seren 🌊

A Self-Refining Decoder-Only LLM Architecture

Seren is a production-grade, HuggingFace-compatible transformer library built around a novel two-pass self-refinement mechanism: the model generates a first response normally, then conditions a second, improved response on a learned projection of its own internal states.

---

Architecture at a glance

Feature	Detail
Attention	Grouped Query Attention (GQA) + Flash Attention 2 via torch.nn.functional.scaled_dot_product_attention
Position	Rotary Position Embedding (RoPE)
Normalisation	Pre-RMSNorm
FFN	SwiGLU (gated MLP)
KV Cache	Fully compatible — no tokens discarded between passes
Gradient Checkpointing	✅
Weight Tying	Optional (tie_word_embeddings)
Device Parallelism	2-GPU via HuggingFace device_map
Self-Refinement	Novel two-pass mechanism (see below)

---

Installation

pip install seren
# or from source:
pip install -e ".[dev]"

Requirements: Python ≥ 3.8, PyTorch ≥ 2.0, Transformers ≥ 4.36

---

Quick start

from seren import SerenConfig, SerenForCausalLM

config = SerenConfig(
    vocab_size=32_000,
    hidden_size=2_048,
    num_hidden_layers=16,
    num_attention_heads=16,
    num_key_value_heads=4,    # GQA: 4 KV heads shared across 16 query heads
    intermediate_size=5_504,
    max_position_embeddings=4_096,
    self_hidden_layer=2,    # depth of refinement MLP (2 linear layers)
    alpha=1.0,                # weight for loss₁ (Pass-1 quality)
    beta=1.0,                 # weight for loss₂ (refinement contrastive)
)

model = SerenForCausalLM(config)
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

---

The Self-Refinement Mechanism

Seren generates every response in two passes over the same sequence:

[System Prompt] + [User] + [Response₁] + [Response₂]

Pass 1 — Baseline generation

The model runs a full forward pass on the entire sequence with no special conditioning. The final-layer hidden states at the RESP1 positions are captured.

Refinement projection

A depth-configurable MLP (self_hidden_layer linear layers, each hidden_size → hidden_size with SiLU between them) followed by RMSNorm processes the Pass-1 hidden states and compresses them via mean-pooling into a single [B, 1, hidden_size] conditioning vector.

Pass 2 — Refined generation

The conditioning vector is added to the token embeddings of the RESP2 positions before the second forward pass. The model then predicts a refined second response conditioned on both the original context and its own internal Pass-1 representation.

Training objective

# loss₁ — forces the first response to be good
loss1 = cross_entropy(pass1_logits[resp1_positions], resp1_labels)

# loss₂ — forces the second response to improve over the first
#   CE(pass1 at resp2) = baseline CE without refinement
#   CE(pass2 at resp2) = CE with refinement → should be lower → loss₂ > 0
loss2 = cross_entropy(pass1_logits[resp2_positions], resp2_labels) \
      - cross_entropy(pass2_logits[resp2_positions], resp2_labels)

total_loss = alpha * loss1 + beta * loss2

---

Training usage

import torch
from seren import SerenForCausalLM, SerenConfig

model = SerenForCausalLM(SerenConfig()).cuda()

# Prepare batch: [SYS+USR+RESP1] and [RESP2]
pass1_ids    = torch.randint(0, 32000, (2, 64)).cuda()   # [B, T1]
pass2_ids    = torch.randint(0, 32000, (2, 32)).cuda()   # [B, T2]

# Labels: -100 for positions we don't want loss on
resp1_labels = pass1_ids.clone()
resp1_labels[:, :10] = -100   # mask system/user turns

resp2_labels = pass2_ids.clone()

outputs = model(
    input_ids=pass1_ids,
    response2_input_ids=pass2_ids,
    labels=resp1_labels,
    response2_labels=resp2_labels,
)

outputs.loss.backward()
print(f"loss={outputs.loss:.4f}  loss1={outputs.loss1:.4f}  loss2={outputs.loss2:.4f}")

---

Inference: two-pass generation

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("...")
model = SerenForCausalLM.from_pretrained("...", device_map="balanced")

context = tokenizer("User: What is 2+2?\nAssistant:", return_tensors="pt").input_ids

# ── Pass 1: generate first response ───────────────────────────────────
out1 = model.generate(
    input_ids=context,
    max_new_tokens=128,
    return_dict_in_generate=True,
    output_hidden_states=True,
)

# Extract last-layer hidden states from the last generated step
pass1_hidden = out1.hidden_states[-1][-1]   # [B, 1, H]

# ── Compute refinement conditioning vector ─────────────────────────────
refinement_ctx = model.compute_refinement_context(pass1_hidden)

# ── Pass 2: generate refined response ─────────────────────────────────
with model.with_refinement(refinement_ctx):
    out2 = model.generate(
        input_ids=out1.sequences,
        max_new_tokens=128,
    )

print(tokenizer.decode(out2[0]))

---

2-GPU model parallelism

model = SerenForCausalLM.from_pretrained(
    "path/to/checkpoint",
    device_map="balanced",   # splits layers evenly across available GPUs
    torch_dtype=torch.bfloat16,
)

---

Component imports

from seren import SerenConfig           # configuration
from seren import SerenForCausalLM     # full model with LM head + self-refinement
from seren import SerenModel           # transformer body only
from seren import SerenDecoderLayer    # single decoder layer
from seren import SerenRMSNorm         # RMS normalisation
from seren import SerenRotaryEmbedding # RoPE
from seren import SerenCausalLMOutputWithPast  # typed output dataclass

---

License

Apache 2.0