File size: 1,640 Bytes
877ce7f
 
 
 
 
 
8b05f99
877ce7f
 
8391f7c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
877ce7f
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
---
license: apache-2.0
language:
- en
base_model:
- FlameF0X/i3-BERT
pipeline_tag: fill-mask
---

# [i3-BERT](https://github.com/FlameF0X/open-i3/tree/main/src/fill-mask/BERT): Hybrid RWKV-Transformer for Efficient Pre-training

A novel hybrid language model architecture combining the efficiency of RWKV's linear attention with the global reasoning capabilities of standard transformers, designed for BERT-style masked language modeling tasks.

## Architecture Overview

**i3-BERT** implements a two-tier architecture:

- **Bottom Layers (Bi-RWKV)**: Process local context efficiently using bidirectional RWKV blocks with O(T) complexity
- **Top Layers (Full Attention)**: Perform global reasoning and long-range dependencies with O(T²) multi-head attention

This design philosophy leverages the strengths of both approaches: RWKV handles syntactic structure and local patterns efficiently, while attention layers enable global information retrieval and complex reasoning.

## Key Features

- **Bidirectional RWKV**: Novel implementation running RWKV in both forward and backward directions for non-causal tasks
- **JIT-Optimized WKV Kernel**: Compiled linear attention mechanism for faster training
- **Hybrid Layer Stack**: Configurable ratio of 4 Bi-RWKV to 4 Attention layers
- **Standard BERT Pre-training**: MLM (Masked Language Modeling) + NSP (Next Sentence Prediction)
- **Streaming Data Pipeline**: Handles large datasets without memory issues
- **116M Parameters**: Educational-scale model suitable for consumer GPUs

Iter 0 | Loss: 11.2089 | MLM: 10.4452 | NSP: 0.7637
...
Iter 4990 | Loss: 0.1881 | MLM: 0.1489 | NSP: 0.0392