Skip to content

Add Arcee squared-ReLU MLP adapter (ArceeForCausalLM)#1478

Merged
jlarson4 merged 4 commits into
TransformerLensOrg:devfrom
pablocs116:arcee-squared-relu-mlp
Jul 2, 2026
Merged

Add Arcee squared-ReLU MLP adapter (ArceeForCausalLM)#1478
jlarson4 merged 4 commits into
TransformerLensOrg:devfrom
pablocs116:arcee-squared-relu-mlp

Conversation

@pablocs116

Copy link
Copy Markdown

Adds TransformerBridge support for Arcee AI's AFM-4.5B (ArceeForCausalLM), a Llama-style dense decoder whose only architectural novelty is an ungated squared-ReLU (ReLU^2) MLP (up_proj -> ReLU^2 -> down_proj) instead of the gated SiLU/GeLU used by Llama.

  • Register the "relu2" activation (relu(x)**2) in SUPPORTED_ACTIVATIONS so the config guard and ActivationFunctionFactory accept HF's hidden_act="relu2".
  • New ArceeArchitectureAdapter: Llama attention (RoPE, RMSNorm, GQA, no QK-norm, no biases) + ungated MLPBridge. Post-activation ReLU^2 neurons are inspectable via blocks.{i}.mlp.hook_post.
  • Register at the four sites (factory, supported_architectures init, model_registry HF_SUPPORTED_ARCHITECTURES + CANONICAL_AUTHORS_BY_ARCH, generate_report descriptions) and add the two checkpoints to the registry.
  • Unit + integration parity tests (fp32 + eager); integration asserts bridge logits match HF and hook_post exposes non-negative d_mlp-width activations.

Closes #1467

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant.

  • New feature (non-breaking change which adds functionality)

Screenshots

Please attach before and after screenshots of the change if applicable.

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

danra and others added 4 commits June 18, 2026 18:04
Adds TransformerBridge support for Arcee AI's AFM-4.5B (ArceeForCausalLM), a
Llama-style dense decoder whose only architectural novelty is an ungated
squared-ReLU (ReLU^2) MLP (up_proj -> ReLU^2 -> down_proj) instead of the gated
SiLU/GeLU used by Llama.

- Register the "relu2" activation (relu(x)**2) in SUPPORTED_ACTIVATIONS so the
  config guard and ActivationFunctionFactory accept HF's hidden_act="relu2".
- New ArceeArchitectureAdapter: Llama attention (RoPE, RMSNorm, GQA, no QK-norm,
  no biases) + ungated MLPBridge. Post-activation ReLU^2 neurons are inspectable
  via blocks.{i}.mlp.hook_post.
- Register at the four sites (factory, supported_architectures __init__,
  model_registry HF_SUPPORTED_ARCHITECTURES + CANONICAL_AUTHORS_BY_ARCH,
  generate_report descriptions) and add the two checkpoints to the registry.
- Unit + integration parity tests (fp32 + eager); integration asserts bridge
  logits match HF and hook_post exposes non-negative d_mlp-width activations.

Closes TransformerLensOrg#1467

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@jlarson4 jlarson4 changed the base branch from main to dev July 2, 2026 15:25
@jlarson4

jlarson4 commented Jul 2, 2026

Copy link
Copy Markdown
Collaborator

Looks great thank you @pablocs116

@jlarson4 jlarson4 merged commit 77a44c0 into TransformerLensOrg:dev Jul 2, 2026
25 checks passed
@pablocs116 pablocs116 deleted the arcee-squared-relu-mlp branch July 2, 2026 15:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Proposal] Add Arcee squared-ReLU sparse-MLP adapter (ArceeForCausalLM)

3 participants