Introduction

This repository contains the core implementation of our proposed method FlashShield. It focuses on the modeling and detection module described in the paper.

To ensure the verifiability of the experimental results, we include the preprocessed training and testing datasets (D1) used in the paper. In addition, the best-performing model weights and hyperparameters are provided to allow for direct validation.

---

Environment Setup

The project was developed and tested with:

• Python 3.11
• PyTorch 1.13.1 (CUDA 11.7)
• Managed under Anaconda

To create the environment:

conda env create -f environment.yml
conda activate flashshield

If you only need CPU testing:

conda install pytorch==1.13.1 cpuonly -c pytorch

---

Configuration

All key parameters and file paths are defined in config.ini. You can modify:

• Model name and architecture
• Input/output paths
• Hyperparameters and thresholds
• Logging or saving options

See Dataset Structure for data format details and Hypergraph Configuration for hyperedge construction and training parameters.

---

Dataset Structure

Overview and Directory Layout

dataset/
├── feature_mapping.json             # Feature name → column index mapping (264 features)
├── train_dataset/                   # Training set (2,440 transactions)
│   ├── hyperGraph-attacks/
│   │   ├── nodes/                   # 122 raw attack transaction files
│   │   └── new_nodes/               # 122 anonymized versions
│   ├── hyperGraph-normalTX/
│   │   ├── nodes/                   # 2,318 raw normal transaction files
│   │   └── new_nodes/               # 2,318 anonymized versions
│   └── processed_data/
│       ├── mix_nodes/               # 2,440 renumbered JSON files (0.json, 1.json, ...)
│       ├── all_nodes.json           # JSONL: 29,494 node feature vectors (264-dim each)
│       ├── file_id_mapping.json     # Transaction ID → tx hash
│       ├── global_id_mapping.json   # Transaction ID → node global IDs
│       ├── global_features.json     # Tx hash → 15 graph-level features
│       └── index_label_mapping.json # Transaction ID → "normal" / "attack"
└── test1_dataset/                   # D1 test set (same structure, 620 transactions)

Each transaction is modeled as a hypergraph in which nodes represent individual token transfers and hyperedges capture semantic relationships among them.

Raw Transaction Data Format

Files under nodes/ are named nodes_{txHash}.json. Each file contains:

{
    "TxHash": "0x0016745693d68d734faa408b94cdf2d6c95f511b50f47b03909dc599c1dd9ff6",
    "NodeNum": 6,
    "Nodes": [
        {
            "Index": 0,
            "TokenName": "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2",
            "Amount": "100000000000000000000",
            "From": "uniswap",
            "To": "0xcE1F4B4F17224ec6df16Eeb1e3e5321c54Ff6EDe"
        },
        {
            "Index": 1,
            "TokenName": "0xD06527D5e56A3495252A528C4987003b712860eE",
            "Amount": "483446999826",
            "From": "cream-finance",
            "To": "0xcE1F4B4F17224ec6df16Eeb1e3e5321c54Ff6EDe"
        }
    ]
}

• TxHash: Transaction hash on the blockchain.
• NodeNum: Number of token transfers (nodes) in this transaction.
• Nodes: Array of token transfers. Each node contains:
  • Index: Sequential index within the transaction.
  • TokenName: ERC-20 token contract address (or token identifier).
  • Amount: Transfer amount as a string in wei.
  • From: Sender address — either a hex string or a recognized DeFi protocol label (e.g., uniswap, aave, compound).
  • To: Receiver address — same format as From.

Anonymization and Normalization

Files under new_nodes/ are derived from nodes/ through the following transformations (performed by node_process.py):

Transformation	Details
Token pseudonymization	Contract addresses are replaced with Token1, Token2, ..., up to Token20.
Address pseudonymization	Unknown addresses are replaced with From1–From100 (senders) or To1–To100 (receivers).
DeFi protocol labels	Recognized protocol names (uniswap, aave, compound, sushiswap, etc.) are preserved as-is.
Amount normalization	log1p(amount) followed by min-max normalization to [0, 1] within each transaction. Single-node transactions receive 0.0.

Important: The pseudonymization is scoped per-transaction. Token1 or From1 in one transaction refers to a different real address than Token1 or From1 in another transaction.

Processed Data Formats

The processed_data/ directory contains the final representations consumed by the model:

File	Format	Description
mix_nodes/	Directory of {id}.json	All transactions from new_nodes/, shuffled and renumbered starting from 0.json. Each file retains the same JSON structure as new_nodes/.
all_nodes.json	JSONL (one JSON object per line)	All nodes concatenated globally. Each line is a 264-dim feature vector: {"1": global_id, "2": index, "3": amount, "4"–"264": one-hot}. Keys are string column indices from feature_mapping.json.
file_id_mapping.json	{"0": "0xhash...", "1": "0xhash...", ...}	Maps sequential transaction ID to the original transaction hash.
global_id_mapping.json	{"0": [{"global_id": 0, "original_index": 0}, ...], ...}	Maps each transaction ID to the list of global node IDs (and their original indices) belonging to that transaction.
global_features.json	{"0xhash...": {"n_nodes": 4, ...}, ...}	Maps transaction hash to 15 graph-level statistical features (see Global Graph-Level Features).
index_label_mapping.json	{"0": "normal", "1": "normal", ..., "8": "attack", ...}	Maps sequential transaction ID to its label ("normal" or "attack").

Node Feature Space (264 Dimensions)

Each node is represented by a 264-dimensional feature vector. The feature layout, defined in dataset/feature_mapping.json, is as follows:

Column Range	Dim	Description
1	1	Global node ID (sequential across all transactions)
2	1	Node index within its transaction
3	1	Normalized transfer amount (log1p + min-max, per-transaction)
4–19	16	TokenName one-hot encoding (Token1–Token15, Token20)
20–43	24	From address — DeFi protocol one-hot (1inch, aave, alchemix-finance, balancer, bancor, bzx, compound, convex-finance, cream-finance, curve-fi, defi, dydx, idle-finance, inverse-finance, kyberswap, maker, pancake-bunny, quickswap, rari-capital, sushiswap, synthetix, uniswap, yearn-finance, yield-farming)
44–143	100	From address — anonymized address one-hot (From1–From100)
144–164	21	To address — DeFi protocol one-hot (same as From, minus alchemix-finance, compound, inverse-finance)
165–264	100	To address — anonymized address one-hot (To1–To100)

The one-hot encoding is generated by sklearn.preprocessing.OneHotEncoder fitted on the new_nodes/ data. Categories not present in a given dataset are encoded as 0.0.

Global Graph-Level Features

In addition to per-node features, 15 graph-level features are computed per transaction by node_process.py. Of these, 13 are used by the model (defined in the G_KEYS constant in train.py and test1.py):

Feature	Description	Used by Model
n_nodes	Number of token transfers in the transaction	Yes
n_addrs	Number of distinct addresses (From + To union)	Yes
n_tokens	Number of distinct tokens	Yes
avg_amount_rel	Mean of normalized amounts	Yes
max_amount_rel	Maximum normalized amount	Yes
std_amount_rel	Standard deviation of normalized amounts	Yes
in_deg_max	Maximum in-degree across all addresses	Yes
in_deg_mean	Mean in-degree across all addresses	Yes
out_deg_max	Maximum out-degree across all addresses	Yes
out_deg_mean	Mean out-degree across all addresses	Yes
amount_gini	Gini coefficient of the normalized amount distribution	Yes
token_entropy	Shannon entropy of the token frequency distribution	Yes
addr_entropy	Shannon entropy of the address frequency distribution	Yes
amount_hhi	Herfindahl–Hirschman Index of amounts	No
cycle_len_min	Shortest directed cycle length (BFS up to depth 6)	No

The 13 active features are standardized using a RobustScaler fitted on the training set and saved in model/global_scaler/.

Dataset Statistics

Split	Normal	Attack	Total	Total Nodes
Train	2,318	122	2,440	29,494
Test (D1)	589	31	620	5,884

The class imbalance ratio is approximately 19:1 (normal:attack) in both splits.

---

Hypergraph Configuration

Hyperedge Construction Heuristics

For each transaction, a hypergraph is constructed where nodes are token transfers and hyperedges represent semantic relationships. Five types of hyperedges are built (implemented in semantic_hyperedge.py):

1. KNN Hyperedges (hyperedges_knn, default: enabled, k=3)

   • For each node, a hyperedge is created containing that node and its k nearest neighbors in the 264-dimensional feature space.
   • Captures nodes with similar behavioral profiles regardless of their position in the transaction graph.

2. Same-Address Hyperedges (hyperedges_same_address, default: enabled, mode="either")

   • Nodes sharing the same From address or the same To address are grouped into one hyperedge.
   • Captures all transfers involving a common address, critical for identifying fund routing patterns.

3. Same-Token Hyperedges (hyperedges_same_token, default: enabled)

   • Nodes with the same TokenName are grouped together.
   • Captures multi-hop token flow through the same ERC-20 token, relevant for identifying arbitrage loops.

4. Temporal Window Hyperedges (hyperedges_temporal_window, default: enabled, window=3)

   • Creates sliding-window hyperedges of consecutive nodes: {i, i+1, ..., i+window−1}.
   • Captures sequential execution order of token transfers within the transaction.

5. Flow Chain Hyperedges (hyperedges_flow_chain, default: enabled)

   • For every address that appears as a To (receiver of node i) and also as a From (sender of node j), nodes i and j are grouped into the same hyperedge.
   • Directly models the transaction flow graph, where funds move from one transfer to another through intermediate addresses.

After construction, all hyperedges are deduplicated and filtered:

• Hyperedges with fewer than 2 nodes are removed.
• Hyperedges exceeding max_size (default 16) are chunked into smaller edges.
• A fallback ensures every node belongs to at least one hyperedge.

Configuration Parameters

All parameters are defined in config.ini.

Hyperedge construction ([hypergraph] section):

Parameter	Type	Default	Description
use_knn	bool	True	Enable KNN-based hyperedges
use_same_address	bool	True	Enable address-based hyperedges
use_same_token	bool	True	Enable token-based hyperedges
use_temporal	bool	True	Enable temporal window hyperedges
use_flow	bool	True	Enable flow chain hyperedges
KNN_NUM	int	3	Number of nearest neighbors for KNN hyperedges
max_size	int	16	Maximum hyperedge size (larger edges are chunked)

Training and inference ([param] section):

Parameter	Type	Default	Description
class_balanced_weight_beta	float	0.9999	Beta parameter for class-balanced loss weighting
use_class_balanced_weight	bool	True	Whether to apply class-balanced weights to the loss
threshold	float	0.7	Initial decision threshold for attack classification
focal_loss_gamma	float	2.0	Focal loss focusing parameter
learning_rate	float	0.0001	Adam optimizer learning rate
search_threshold_metric	str	"fp_rate"	Metric for automatic threshold search ("fp_rate", "f1", etc.)
search_threshold_target	float	0.003	Target value for the threshold search metric
if_use_alert	bool	False	Use alert-rate-based dynamic threshold instead of fixed threshold
alert_rate	float	0.025	Target alert rate (fraction of transactions flagged)

Model Architecture Summary

The model processes each transaction independently:

1. Input: 264-dim node features + 13-dim global features per transaction.
2. Hypergraph construction: The five heuristic types above produce a set of hyperedges per transaction via dhg.Hypergraph.
3. HGNN backbone: Two HGNNConv layers (264 → 256 → 64). Each layer fuses three feature views with learnable scalar weights: spectral (HGNN smoothing), spatial (vertex-to-vertex mean aggregation), and original features: α · spectral + β · spatial + γ · original.
4. Readout: Mean-pooling of node embeddings per transaction, then concatenation with 13-dim scaled global features → Linear(77, 2).
5. Training: Focal loss (γ=2.0) with class-balanced weights (β=0.9999), Adam optimizer (lr=0.0001), StepLR scheduler.
6. Post-training calibration: Temperature scaling on the validation set, followed by threshold search optimizing for fp_rate ≤ 0.003.

---

Time-Based Split Experiment

In addition to the default random split (D1), we provide a time-based split experiment where transactions are ordered chronologically and divided into train/test, better reflecting real-world deployment scenarios where the model encounters future attacks it has never seen.

Split Strategies

Two strategies are implemented:

Strategy	Description	Train Attack%	Test Attack%
Global sorted (v1)	All transactions sorted by timestamp, first 80% → train, last 20% → test	1.5%	16.1%
Per-class sorted (v2)	Normal and attack transactions sorted separately, each split 80/20, then combined	4.4%	4.4%

v1 strictly respects temporal ordering but results in highly imbalanced class distributions across splits (normal transactions are older on average). v2 ensures balanced class ratios in both splits at the cost of looser temporal ordering.

Pipeline

# Step 1: Fetch block timestamps from on-chain (requires RPC access)
python scripts/fetch_timestamps.py

# Step 2: Build time-based split
python scripts/build_time_split.py              # v1: global sorted
python scripts/build_time_split_v2.py           # v2: per-class sorted

# Step 3: Process data
python data_process/mix_and_map_time.py train_time
python data_process/mix_and_map_time.py test_time

# Step 4: Extract features and labels
cd data_process
python extract_features.py train_time
python extract_labels.py train_time
python extract_features.py test_time
python extract_labels.py test_time
cd ..

# Step 5: Train and evaluate
python train_time.py
python test_time.py

For the v2 split, replace the config section names with _v2 suffixes:

python data_process/mix_and_map_time.py train_time_v2 --tx-lists-dir dataset/tx_lists_time_v2 --csv-prefix train
python data_process/mix_and_map_time.py test_time_v2 --tx-lists-dir dataset/tx_lists_time_v2 --csv-prefix test
cd data_process
python extract_features.py train_time_v2
python extract_features.py test_time_v2
cd ..
python train_time.py --config-section train_time_v2
python test_time.py --config-section test_time_v2

Key Scripts

File	Description
scripts/fetch_timestamps.py	Batch-fetches block timestamps via JSON-RPC for all transactions
scripts/build_time_split.py	v1: global chronological 80/20 split
scripts/build_time_split_v2.py	v2: per-class chronological 80/20 split
data_process/mix_and_map_time.py	Time-ordered mix-and-map (replaces random shuffle)
train_time.py	Training with temporal train/val split
test_time.py	Evaluation on time-held-out test set
model_common.py	Shared model classes (HGNN, GlobalScaler, FocalLoss, etc.)

---

Quick Start

After setting up the environment, simply run:

python test1.py

The script will:

• Load the pretrained model and D1 test set
• Perform detection on the test data
• Print evaluation metrics on screen
• Save results and logs into the output folder

To use a specific model checkpoint:

python test1.py --model_state_path path/to/model.pht

---

Notes

• The provided datasets are preprocessed subsets (D1) provided for reproducibility.
• The pretrained model weights and hyperparameters correspond to the best-performing configuration reported in the paper.
• Additional datasets and full preprocessing code may be released in future updates.

---

Usage Notice

If you use this code or data in your research, please cite our paper.

This repository is released for academic and research purposes. The included data and models are provided as-is for reproducibility of the experimental results.