PoHao/opt2.7-fc1-predictor

OPT-2.7B fc1 Activation Predictors

Small predictor models for facebook/opt-2.7b. Each predictor looks at the input to one decoder layer's fc1 and predicts which output neurons will be activated, so those neurons can be skipped during inference (sparse inference).

One predictor per decoder layer (32 total): model_0.pt ... model_31.pt.

Pairs with the fine-tuned model at PoHao/opt2.7-predictor-guided-sft.

Architecture

import torch.nn as nn

class PredictorMLP(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim=512):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim, bias=False)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return x

• input_dim: OPT-2.7B hidden size
• output_dim: OPT-2.7B fc1 output size
• hidden_dim: 512

Output is a value in [0, 1] per neuron; threshold (default 0.5) to get a binary mask.

Loading

import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM

base_model = AutoModelForCausalLM.from_pretrained("facebook/opt-2.7b")
input_dim = base_model.model.decoder.layers[0].fc1.in_features
output_dim = base_model.model.decoder.layers[0].fc1.out_features

def load_predictor(layer_idx, repo_id="PoHao/opt2.7-fc1-predictor", hidden_dim=512, device="cpu"):
    path = hf_hub_download(repo_id=repo_id, filename=f"model_{layer_idx}.pt")
    predictor = PredictorMLP(input_dim, output_dim, hidden_dim).to(device)
    predictor.load_state_dict(torch.load(path, map_location=device))
    predictor.eval()
    return predictor

predictor_layer0 = load_predictor(0)

Using the mask

threshold = 0.5

with torch.no_grad():
    prob = predictor_layer0(hidden_state)   # hidden_state: input to layer 0's fc1
    mask = (prob > threshold).float()
    fc1_output = fc1_output * mask