# Parameter Golf v2

Techniques from the [OpenAI Parameter Golf](https://github.com/openai/parameter-golf) competition and [MicroGPT](https://github.com/cagataycali/strands-microgpt) applied to Neon's action heads.

---

## The Key Insight

```mermaid
graph LR
    BB["Frozen Backbone<br/>3-7B params"] --> FUS["Fusion"] --> AH["Action Heads<br/>~2M params"]

    style BB fill:#1565c0,color:#fff
    style AH fill:#e65100,color:#fff
```

The backbone is frozen. **All learning happens in ~2M parameters.** These action heads *are* our small model — every Parameter Golf trick about making small models punch above their weight applies directly.

---

## Six Techniques

### 1. ReLU² Activation

```python
class ReLUSquared(nn.Module):
    """max(0, x)² — smoother than GELU, cheaper to compute."""
    def forward(self, x):
        return F.relu(x).square()
```

**Why:** Gradient is `2·max(0,x)` — continuous at zero (unlike ReLU's discontinuity). For action heads, this produces smoother trajectories without the computational cost of GELU's `exp()`.

### 2. RMSNorm

```python
class RMSNorm(nn.Module):
    """scale = (mean(x²) + ε)^(-0.5) — no mean subtraction."""
    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
```

**Why:** LayerNorm subtracts the mean AND divides by std — two operations. RMSNorm just divides by RMS. Same stabilization, fewer ops. From MicroGPT's `_rmsnorm()`.

### 3. Learnable Residual Scales

```python
self.residual_scales = nn.ParameterList([nn.Parameter(torch.ones(1)) for _ in range(n)])

# In forward:
h = h + scale * h_in  # scale is learned, starts at 1.0
```

**Why:** Frozen backbone features can dominate the learned action features early in training. The learned α lets the network decide how much residual to keep per layer.

### 4. U-Net Skip Connections

```python
# Layer 0 output saved as 'skip'
# Last hidden layer receives: concat(h, skip)
h = torch.cat([h, skip], dim=-1)
```

**Why:** In 3-4 layer MLPs, the middle layers can become a gradient bottleneck. Skipping from layer 0 to the last layer creates a gradient highway. Only activates for `num_layers >= 3`.

### 5. Logit Soft-Capping

```python
def soft_cap(x, cap=1.0):
    """cap · tanh(x / cap) — smoother than hard Tanh."""
    return cap * torch.tanh(x / cap)
```

```mermaid
graph LR
    subgraph "Hard Tanh"
        HT["Output hits wall at ±1<br/>Zero gradient beyond"]
    end
    subgraph "Soft-Cap"
        SC["Output approaches ±1 asymptotically<br/>Always has nonzero gradient"]
    end

    style SC fill:#1b5e20,color:#fff
    style HT fill:#b71c1c,color:#fff
```

**Why:** Hard Tanh kills gradients when actions are near the boundary (±1). Soft-capping maintains gradient flow, letting the optimizer fine-tune actions at the extremes.

### 6. Optimizer Tuning

| Setting | Standard | Parameter Golf | Why |
|---|---|---|---|
| `adam_beta1` | 0.9 | **0.85** | Faster adaptation — action distributions shift during training |
| `adam_beta2` | 0.999 | **0.99** | Less momentum lag from MicroGPT |
| `max_grad_norm` | 1.0 | **0.3** | Small heads diverge easier — tighter clip |

---

## Configuration

All v2 features are **on by default** but individually togglable:

```python
from neon.model.action_heads import ActionHeadConfig

# Full v2 (default)
config = ActionHeadConfig(
    use_relu_squared=True,      # ReLU² instead of GELU
    use_rmsnorm=True,           # RMSNorm before each layer
    use_residual_scale=True,    # Learned α for residuals
    use_skip_connections=True,  # U-Net skip layer 0 → last
    soft_cap_value=1.0,         # Soft-cap output range
)

# Disable v2 (legacy mode)
config = ActionHeadConfig(
    use_relu_squared=False,
    use_rmsnorm=False,
    use_residual_scale=False,
    use_skip_connections=False,
    soft_cap_value=0.0,         # Hard Tanh
)
```

---

## Origin

These techniques were discovered by analyzing winning entries in OpenAI's Parameter Golf competition, where contestants minimized perplexity per parameter byte. The pure-Python implementation lives in [MicroGPT's engine.py](https://github.com/cagataycali/strands-microgpt) — we ported the concepts to PyTorch for Neon's action decoders.
