| --- |
| license: apache-2.0 |
| tags: |
| - jax |
| - flax |
| - research |
| - orthogonal-embeddings |
| - walsh-hadamard |
| - gram-schmidt |
| - efficient-transformers |
| datasets: |
| - roneneldan/TinyStories |
| - HuggingFaceFW/fineweb-edu |
| - HuggingFaceTB/smollm-corpus |
| - HuggingFaceTB/dclm-edu |
| --- |
| |
| # Supra-Mini com Difusão Isométrica de Hadamard e Anti-Resíduo de Gram-Schmidt Dinâmico |
|
|
| Este repositório documenta a implementação, teoria e checkpoints de uma arquitetura Transformer compacta (**812.800 parâmetros**) projetada para mitigar dois gargalos teóricos fundamentais em modelos de linguagem: |
| 1. **O custo de "apagar" o token de entrada no barramento residual (*Residual Memory Burden*).** |
| 2. **O desperdício paramétrico de projeções de posto completo em cada subcamada incremental.** |
|
|
| --- |
|
|
| ## 1. Fundamentos da Arquitetura |
|
|
| ``` |
| [Token Entrada x_0] (128d) |
| │ |
| ▼ |
| ┌─────────────────────────────────────────────────────────────────┐ |
| │ Subcamada l (Atenção ou FFN): │ |
| │ 1. Extração não-linear em subespaço local de baixo rank (64d) │ |
| │ 2. Padding com 64 zeros (64 -> 128) │ |
| │ 3. Rotação pseudoaleatória fixa Π_l (Quebra simetria diádica) │ |
| │ 4. Difusão isométrica máxima: Multiplicação por Hadamard (128) │ |
| │ 5. Soma linear no barramento residual: x = x + Δx_l │ |
| └─────────────────────────────────────────────────────────────────┘ |
| │ (12 subcamadas com subespaços mutualmente incoerentes) |
| ▼ |
| [Fluxo Residual Acumulado de Posto Completo 128d] |
| │ |
| ▼ |
| ┌─────────────────────────────────────────────────────────────────┐ |
| │ Cabeça de Decodificação: │ |
| │ Filtro Dinâmico de Gram-Schmidt com Gate Dependente de h: │ |
| │ g(h) = sigmoid(W_gate @ h_final + b) │ |
| │ h_orth = h_final - g(h) ⊙ proj_{x_0}(h_final) │ |
| └─────────────────────────────────────────────────────────────────┘ |
| │ |
| ▼ |
| [Tied Softmax Head: W_embed^T @ h_orth] |
| ``` |
|
|
| ### Principais Inovações Teóricas: |
| 1. **Atualizações de Baixo Rank com Difusão Isométrica de Hadamard:** |
| * Cada subcamada projeta apenas em $64$ dimensões ativas e preenche com $64$ zeros. |
| * A matriz ortonormalizada de Sylvester-Hadamard ($H^T H = I$, $|H_{ij}| = 1/\sqrt{128}$) espalha a energia de forma equiprovável sobre todas as 128 coordenadas sem alterar a norma $\ell_2$. |
| * As permutações $\{\Pi_1, \dots, \Pi_{12}\}$ garantem que os subespaços de atualização das 12 subcamadas tenham distância Grassmanniana máxima, permitindo cobrir todo o espaço $\mathbb{R}^{128}$ no acumulador. |
| 2. **Projetor de Gram-Schmidt com Válvula de Escape Dinâmica:** |
| * Modela algebricamente a subtração da componente colinear a $x_0$, desonerando as camadas ocultas de aprenderem interferência destrutiva. |
| * O gate $g(h) \in (0, 1)$ impede o colapso de repetição: quando a sintaxe exige repetição legítima de tokens (ex: `--`, hifens de lista `- `, quebras de linha `\n`), o gate desce para $\approx 0.18 - 0.25$, preservando $x_0$. |
|
|
| --- |
|
|
| ## 2. Checkpoints Disponíveis no Repositório |
|
|
| Os pesos estão serializados no formato Flax (`.msgpack`) no diretório `checkpoints/`: |
|
|
| | Checkpoint | Tokens Acumulados | Corpus de Treino | Função de Perda | Observações | |
| | :--- | :--- | :--- | :--- | :--- | |
| | `supra_mini_tied_gs_params.msgpack` | **10M** | TinyStories | Hard Cross-Entropy | Loss: **2.7381**. Throughput: 246k tok/s. | |
| | `supra_mini_20m_general_mix.msgpack` | **30M** | Mix Web (50% FineWeb, 30% Cosmo, 10% DCLM, 10% Tiny) | Hard Cross-Entropy | Loss: **4.7582**. Domínio de sintaxe e markdown. | |
| | `supra_mini_70m_general_mix.msgpack` | **80M** | Mix Web Geral | Hard CE (Anneal $5e-4 \to 5e-5$) | Loss estabilizado; limpeza de atratores espúrios. | |
| | `supra_mini_75m_self_distill.msgpack` | **85M** | Mix Web Geral | $0.65 \text{ OneHot} + 0.35 \text{ Top-5 Soft}$ | Queda abrupta de entropia ($-0.85$ nats). | |
| | `supra_mini_90m_self_distill.msgpack` | **100M** | Mix Web Geral | $0.65 \text{ OneHot} + 0.35 \text{ Top-5 Soft}$ | **Ponto fixo assintótico comprovado** ($\Delta H \approx -0.04$ nats). | |
|
|
| --- |
|
|
| ## 3. Como Carregar e Usar em Python / JAX |
|
|
| ```python |
| from pathlib import Path |
| import jax |
| import jax.numpy as jnp |
| import flax |
| from transformers import AutoTokenizer |
| from modeling_supra_gs import SupraMiniWithGate, generate_hadamard_matrix, generate_sublayer_permutations |
| |
| # 1. Tokenizer e Matrizes Estruturadas |
| tokenizer = AutoTokenizer.from_pretrained("SupraLabs/Supra-Mini-v6-1M") |
| H_128 = generate_hadamard_matrix(128) |
| perms = generate_sublayer_permutations(num_layers=6, dim=128, seed=2026) |
| |
| # 2. Inicialização do Modelo |
| model = SupraMiniWithGate( |
| vocab_size=tokenizer.vocab_size, |
| d_model=128, |
| num_layers=6, |
| num_heads=4, |
| d_head=16, |
| max_len=256 |
| ) |
| |
| # 3. Carregar Pesos do Checkpoint (Exemplo: 90M) |
| dummy_input = jnp.zeros((1, 256), dtype=jnp.int32) |
| template_params = model.init(jax.random.PRNGKey(0), dummy_input, perms, H_128)["params"] |
| |
| ckpt_file = "checkpoints/supra_mini_90m_self_distill.msgpack" |
| with open(ckpt_file, "rb") as f: |
| params = flax.serialization.from_bytes(template_params, f.read()) |
| |
| # 4. Inferência |
| prompt = "Photosynthesis is a biological process where plants use sunlight to" |
| tokens = tokenizer.encode(prompt) |
| tok_tensor = jnp.zeros((1, 256), dtype=jnp.int32).at[0, :len(tokens)].set(jnp.array(tokens)) |
| |
| logits, gate_h, proj_x_in, _ = model.apply({"params": params}, tok_tensor, perms, H_128) |
| print(f"Logits shape: {logits.shape} | Gate médio: {jnp.mean(gate_h[0, :len(tokens)]):.4f}") |
| ``` |
|
|
| --- |
|
|
| ## 4. Hardware e Throughput |
|
|
| * **Acelerador:** Google Cloud TPU v5e-1 core. |
| * **Velocidade Média:** **$186.000$ a $246.000$ tokens/segundo**. |
| * **Tempo Total de Computação:** $< 7$ minutos para executar 100 milhões de tokens de treinamento e avaliação. |
