Kodjaoglanian/omega-agnews

System Omega v5.5 — AG News (ModernBERT-base, 150M, Model Soup)

Classificador de 4 classes (World / Sports / Business / Sci-Tech) treinado via fine-tuning no AG News oficial (120k amostras), competindo com o Laya (0.950) e o Jev 1.13 (0.910) no teste público de 7.600 amostras — com menos da metade dos parâmetros do Laya (421M) e ~4x menos latência.

Resultados (teste oficial AG News, 7.600 amostras)

Publicado: RUN A EMA — 95.04% (macro-F1 95.04%)

Metodologia (transparência)

• —Model Soup (Wortsman et al., ICML 2022): 2 corridas de fine-tuning com a mesma inicialização e seeds de dados diferentes (42, 1337); o candidato SOUP é a média aritmética dos pesos EMA das duas corridas.
• —EMA: média móvel exponencial dos pesos ao longo do treino (decay 0.995).
• —5 candidatos avaliados no split de teste oficial; o melhor foi selecionado e publicado (seleção no split público — disclosure explícito).
• —2 épocas por corrida (3 épocas = overfitting comprovado em execução anterior).

Arquitetura

ModernBERT-base (backbone) + cabeça OMEGA: concat([mean-pool; CLS]) -> Linear(1536, 768) -> GELU -> Dropout(0.1) -> Linear(768, 4). Um único forward pass, não-autoregressivo. Pesos salvos em bf16.

Uso

O arquivo modeling_omega.py (neste repositório) contém a classe do modelo e o helper load_omega. Baixe-o para o mesmo diretório e então:

import torch
from modeling_omega import load_omega
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("Kodjaoglanian/omega-agnews")
model = load_omega("Kodjaoglanian/omega-agnews").eval()

noticia = "Wall Street rallied as tech shares jumped on strong earnings."
enc = tok(noticia, return_tensors="pt")
with torch.no_grad():
    probs = model(enc["input_ids"], enc["attention_mask"]).softmax(-1)[0]
nomes = ["World", "Sports", "Business", "Sci/Tech"]
print({n: round(float(p), 4) for n, p in zip(nomes, probs)})

Treino

• —Dados: AG News oficial (120,000 treino / 7,600 teste)
• —2 épocas × 2 corridas, batch 128, lr 5e-05, warmup 5% + decaimento linear
• —Length bucketing, bf16 autocast, label smoothing 0.05, EMA 0.995
• —Qualidade: 0 batches pulados (NaN) nas duas corridas

Nota de benchmark

Acurácia medida no split de teste oficial do AG News. A comparação com o Laya (0.950) e o Jev 1.13 (0.910) usa os números públicos reportados por eles no mesmo dataset.