brainrot-partition-BTRMplus

BTRM+ (Bradley-Terry Reward Model Plus)

Multi-head reward models for corpus membership and structural genre classification. Trained on situated dialogue from video games and synthetic settings.

Models in This Repository

Training Evolution

Phase 1: Frozen Probes (LoRA)

• Quick validation that Bradley-Terry loss works
• Base transformer frozen, only adapter + BTRM heads trained
• Higher logsquare (0.1) = stronger regularization toward unit logits
• Result: Loss converges, but limited expressivity

Phase 2: Full Fine-Tuning

• Unfroze base transformer for end-to-end training
• Lower logsquare (0.01) = allows larger logit magnitudes
• Added synthetic corpora + structural genre heads
• Result: 2x more weight drift, better discrimination

Weight Drift Analysis

Post-training comparison against original pre-trained weights:

Frozen (LoRA) Models: Zero drift on base transformer

qwen_2head_probe:  0.00 L2 (472M params unchanged)
gemma_2head_probe: 0.00 L2 (253M params unchanged)

Full Fine-Tuned Model: Significant drift, especially in MLP layers

gemma_9head_btrm: 15.53 L2 total (268M params)
  - MLP:       11.20 L2 (3.26% relative)
  - Embedding:  7.94 L2 (1.60% relative)
  - Attention:  7.26 L2 (2.07% relative)
  - Norm:       0.01 L2 (0.00% relative)

Top drifting layers are MLP down_proj weights (up to 15.7% relative change).

Head Types

Corpus Membership (6 heads in 9-head model)

Score whether text belongs to a specific narrative setting:

Structural Genre (3 heads, 9-head model only)

Score text format/style:

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load 9-head model (full fine-tuned)
model = AutoModelForCausalLM.from_pretrained(
    "SQCU/brainrot-partition-BTRMplus",
    subfolder="gemma_9head_btrm/base_model",
    torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained(
    "SQCU/brainrot-partition-BTRMplus",
    subfolder="gemma_9head_btrm/base_model",
)

# Load BTRM heads
from huggingface_hub import hf_hub_download
btrm_path = hf_hub_download(
    "SQCU/brainrot-partition-BTRMplus",
    "gemma_9head_btrm/btrm_heads.pt"
)
btrm_state = torch.load(btrm_path)
# btrm_state["btrm_state_dict"] contains the head weights
# btrm_state["head_names"] = ["skyrim", "oblivion", "fonv", ...]

Training Data

• Reference: Oblivion, Fallout NV, Skyrim dialogue with emotion annotations
• Synthetic: Gallia v9, Marmotte v6, Sanguo v1 (structural translation pipeline)
• Negatives: Cross-corpus soft negatives, Wattpad, FineWeb, WikiText

Architecture

Input Text
    ↓
[Gemma-3 270M Transformer] ← frozen (probes) or fine-tuned (9-head)
    ↓
Last Hidden State (mean pooled)
    ↓
[RMSNorm → Linear(hidden → N_heads)]
    ↓
Per-head logits (soft tanh capped at ±10)

Loss: log(sigmoid(pos - neg)) + logsquare regularization on logit magnitudes.

Observations

1. Reference corpora discriminate better than synthetic (skyrim/oblivion heads accurate, gallia/sanguo confused)
2. Structural heads work excellently - prose vs dialogue vs aesop cleanly separated
3. Full fine-tuning helps - 9-head model achieves lower loss than frozen probes
4. MLP layers adapt most - down_proj weights show highest relative drift

License

Base model weights: Google Gemma License / Qwen License Training data: Bethesda game dialogue (fair use for research), synthetic generation