AI

A Webshell Detection Engine Based on the GPT Architecture

2026年7月7日
Hx Studio
#AI#GPT#PyTorch#Webshell#Machine Learning
A Webshell Detection Engine Based on the GPT Architecture

Model Architecture

The model uses a GPT-3-style decoder-only architecture.

I trained it from scratch. It has 9.98 million parameters in total and was trained with MPS on a Mac with an M1 chip.

Here is the model code:

import torch
import torch.nn as nn
from tokenizers import Tokenizer
import config
import os

class GPTBlock(nn.Module):
    def __init__(self, embed_size, num_heads):
        super().__init__()
        self.ln1 = nn.LayerNorm(embed_size)
        self.attention = nn.MultiheadAttention(
            embed_dim=embed_size, num_heads=num_heads, batch_first=True
        )
        self.ln2 = nn.LayerNorm(embed_size)
        self.ffn = nn.Sequential(
            nn.Linear(embed_size, embed_size * config.FFN_SIZE),  # Expand dimensions
            nn.GELU(),  # Nonlinear activation, smoother than ReLU
            nn.Linear(embed_size * config.FFN_SIZE, embed_size),  # Reduce dimensions
        )

    def forward(self, x, mask=None):
        # 1. Attention and residual connection
        # PyTorch MultiheadAttention returns (output, attention weights).
        # We only need the output here.
        # is_causal=True works with the mask to prevent the model from
        # looking at future tokens.
        attn_output, _ = self.attention(
            x,
            x,
            x,
            attn_mask=mask,
            is_causal=True,
            need_weights=False,
        )
        x = self.ln1(x + attn_output)

        # 2. FFN and residual connection
        ffn_output = self.ffn(x)
        x = self.ln2(x + ffn_output)

        return x


class SimpleGPT(nn.Module):
    def __init__(self, vocab_size, embed_size, max_seq_length, num_heads, num_layers):
        super().__init__()
        # 1. Token and positional embeddings
        # PyTorch nn.Embedding includes learnable weights.
        self.token_embedding = nn.Embedding(vocab_size, embed_size)
        self.position_embedding = nn.Embedding(max_seq_length, embed_size)

        # 2. Stack multiple Transformer blocks
        # nn.ModuleList ensures that PyTorch tracks their parameters.
        self.blocks = nn.ModuleList(
            [GPTBlock(embed_size, num_heads) for _ in range(num_layers)]
        )

        # 3. Language model head
        self.ln_final = nn.LayerNorm(embed_size)
        # Project back to the vocabulary size. bias=False is a common
        # convention that also saves parameters.
        self.lm_head = nn.Linear(embed_size, vocab_size, bias=False)

    def forward(self, input_ids):
        batch_size, seq_length = input_ids.size()
        device = input_ids.device

        # Generate position indices [0, 1, ..., seq_length - 1]
        positions = torch.arange(
            0, seq_length, dtype=torch.long, device=device
        ).unsqueeze(0)

        # Combine token meaning and position
        x = self.token_embedding(input_ids) + self.position_embedding(positions)

        # Generate the causal mask
        mask = nn.Transformer.generate_square_subsequent_mask(seq_length, device=device)

        # Pass the data through every Transformer block
        for block in self.blocks:
            x = block(x, mask=mask)

        # Apply final normalization
        x = self.ln_final(x)

        # Convert the output into vocabulary scores through the LM head
        # Shape: (Batch_Size, Seq_Len, Vocab_Size)
        logits = self.lm_head(x)

        return logits

The key parameters are:

EMBED_SIZE = 256  # 256-dimensional features
MAX_LEN = 1024 * 4  # 4K context window
VOCAB_SIZE = 1024 * 8
NUM_HEADS = 8  # 8 attention heads (256 / 8 = 32 dimensions per head)
NUM_LAYERS = 6  # 6 stacked Transformer blocks
FFN_SIZE = 4

Building the Vocabulary

trainer = BpeTrainer(
    vocab_size=config.VOCAB_SIZE,
    special_tokens=["[UNK]", "[PAD]", "[BOS]", "[EOS]", "[black]", "[white]"],
)

Pretraining

The pretraining data consisted of WordPress, DedeCMS, and webshell code.

Pretraining dataset statistics

Number of PHP files: 2,445
Total raw tokens: 16,139,532
Tokens actually used for training: 4,480,725
Files truncated by MAX_LEN: 593
Average raw tokens per file: 6,601.04
Average training tokens per file: 1,832.61

Testing the Pretrained Model

prompt = """<?php eval($_POST['"""

The model learned a lot of things that were not particularly useful.

Pretrained model output

Supervised Fine-Tuning

Each sample was structured as:

[bos_id] + code_ids + [eos_id]

I added either a [black] or [white] special token after [eos_id] to label the sample.

White samples for SFT:

White SFT samples

Black samples for SFT came from:

After the first training run, the model had a serious false-positive problem. It even classified <?php echo $admin; ?> as a webshell.

The likely reason was that black samples in the SFT dataset were consistently shorter. The model learned the wrong feature and used code length as a signal.

I later introduced injection-based data augmentation:

  1. Randomly extract fragments from white samples and use them as short white samples.
  2. Remove the <?php ?> tags from black samples, inject the remaining code after a ; in white samples, and use the results as black samples.
  3. Balance the number of samples in each group.
🔍 SFT dataset loaded: 780 original webshells, 780 injection-augmented
webshells, 780 normal samples, and 780 short normal samples
(2,098 original normal samples and 780 original webshells).

Evaluation

Evaluation results

Scaling Up Pretraining

Next, I wanted to explore a question that had interested me for a long time: if the SFT process remains completely unchanged, does increasing the amount of pretraining improve the final result?

Based on a common rule of thumb, the optimal number of pretraining tokens is roughly 20 times the model’s parameter count.

This model has 10 million parameters, so the target is about 200 million training tokens.

I collected approximately 180 million PHP tokens from GitHub repositories. About 100 million tokens were actually used in each training epoch, so two epochs came to roughly 200 million tokens.

Large-scale pretraining dataset

I first trained the model on my Mac. After about five to eight hours, it had completed only 1/20 of the training. A full run would have taken anywhere from several days to a week.

I had no choice but to rent a server with an NVIDIA RTX 5090, which completed the training in a little over two hours.

RTX 5090 training progress

RTX 5090 training metrics

Here was the output after the first pretraining batch:

Output after the first pretraining batch

And here was the output after the second batch:

Output after the second pretraining batch

The model had a serious repetition problem.

The false-positive rate also increased substantially after SFT.

Higher false-positive rate after SFT

I suspected this was because I had included webshell code in the pretraining data, and webshell code contains a great deal of repetition. I removed the webshell code and trained the model again.

Batch 1:

Batch 1 without webshell pretraining data

Batch 2:

Batch 2 without webshell pretraining data

SFT results based on Batch 1:

SFT results based on Batch 1

SFT results based on Batch 2:

SFT results based on Batch 2

I also tried SFT alignment without any pretraining.

SFT Batch 1:

SFT Batch 1 without pretraining

SFT Batch 2:

SFT Batch 2 without pretraining

SFT Batch 3:

SFT Batch 3 without pretraining

Based on the samples I observed, my view is that a model with this few parameters cannot fully learn PHP syntax even after pretraining. As a result, pretraining does not provide much help with webshell detection later on. I therefore decided to use SFT directly in subsequent training runs and skip pretraining.

How Model Parameters Affect the Results

Using several data-engineering techniques, I collected all the webshell repositories I could find on GitHub and expanded the number of white samples. This reduced both the false-positive and false-negative rates.

The baseline was a 10-million-parameter model:

embed_size = 256

10-million-parameter baseline

For the 30-million-parameter model:

embed_size = 512

30-million-parameter model with a larger embedding size

Increasing embed_size added 20 million parameters, but did not significantly change the results.

I then tried increasing the feed-forward network size:

FFN_SIZE = 16

This produced a 20-million-parameter model.

20-million-parameter model with FFN_SIZE 16

I also tried:

FFN_SIZE = 32

This produced a 30-million-parameter model.

30-million-parameter model with FFN_SIZE 32

Even with a much larger FFN, the model’s ability to detect webshells did not improve significantly.

I think the main reason is that there are still too few black samples. Even after collecting and deduplicating every black sample I could find on GitHub, there were fewer than 3,000, and a small portion of them were noisy samples that were not actually webshells. High-quality black samples are expensive to obtain. In Alibaba’s FuMo competition, for example, black samples were reportedly collected at a cost of more than RMB 500 each.

Because there are so few black samples, I do not think the current dataset has reached the upper limit of the model’s capacity. With such a limited amount of data, it is difficult to balance the false-positive and false-negative rates. Reducing the false-negative rate inevitably causes a substantial increase in false positives.

Finally, I deployed the model to my VPS. You can try WebShellGPT here:

https://webshell.hx-std.com/