PSSA是一個非Transformer的小型語言模型。它透過一個遞迴狀態空間層一次讀取一個詞元,並保有一個可查詢的情節記憶庫,且在運行時會重寫部分自身權重。此模型完全以Rust從零開始編寫,沒有使用PyTorch、TensorFlow或任何機器學習框架。

在相同參數數量和語料庫條件下,PSSA的學習速度比Transformer快,且在相同CPU上生成文本速度約快十二倍。

這並非為了速度分數,也不是因為語言使架構更好。PSSA需要每個詞元的權重更新、一個在前向傳播中寫入的記憶庫,以及一條所有批次核函數都能微分的標量參考路徑。在自動微分框架中表達這些需求會與框架產生衝突,因此線性代數部分直接手寫,這使得可塑部分簡單且梯度可與參考值校驗至約3e-8。此處的重點是架構本身,實作語言只是細節,歡迎Python版本移植。

Transformer會對上下文中每對詞元進行評分,因此每步的計算成本隨序列長度平方增長,且每步都會重新讀取整個上下文。PSSA則攜帶一個固定大小的狀態,沿序列單向左到右傳遞,並透過記憶庫查詢取代重讀上下文,成本隨長度線性增長。

每個詞元通過一層PSSA:選擇性狀態空間遞迴、在雙曲空間中從情節記憶庫有界讀取、一個學習的門控決定讀取結果進入殘差流的比例,以及一個SiLU多層感知器。預設為256通道,16狀態每通道,以及一個rank-16適配器。

令x為層正規化後的詞元嵌入。從詞元本身讀取三個投影,這使得遞迴是選擇性的而非固定的。

轉移矩陣為對角矩陣,每對(通道, 狀態)有一個速率,透過設計保持為負值,避免遞迴發散。

以步長delta離散化該連續系統,得到每詞元的更新。h在詞元間及訓練時跨區塊傳遞。

A_raw初始化時,每通道的16個速率分布在從1.5到200詞元的對數間隔時間尺度tau上,類似HiPPO初始化。單一通道因此同時持有最近兩個詞元與最近兩百個詞元的資訊,訓練過程會調整這些時間尺度,而非從零發現。

此層的一半為選擇性對角SSM,無新穎性,與S4和Mamba屬同一系列,採先標量後向傳播逐項檢查的寫法。

PSSA特有的是對y的處理。查詢由當前詞元和當前狀態組成,檢索條件不僅依賴詞元,也依賴遞迴狀態。

讀取限制為四個槽位,透過雙曲距離的softmax加權,溫度為tau_mem。

雙曲距離隨球面邊界增大,使得一般上下文槽與特定情節槽保持可分離,且讀取範圍不擴大。四槽為固定成本,與記憶庫大小無關。

讀取結果不無條件加入流中。學習的每通道門控決定讀取量,並搭配低秩SiLU適配器執行目標更新。

寫入是此架構稱為可塑性的原因。當輸入狀態對記憶庫為新穎時插入槽位,每槽位有一個不應期計數器限制覆寫頻率,且快速可塑更新透過封閉形式的嶺迴歸折回基礎轉移矩陣,而非永久存於外部存儲。

不應期計數器防止矛盾更新流抹除已穩定的槽位,整合步驟則避免記憶庫成為唯一長距結構存儲地。

遞迴為標準選擇性SSM機制。創新點為基於遞迴狀態的雙曲有界讀取、新穎與不應期寫入規則,以及快速權重到轉移矩陣的嶺迴歸整合。所有實作均對標量參考路徑微分,批次與平行實作每次提交皆校驗,最大梯度誤差約3e-8。

兩模型在相同語料、分詞器、優化器排程、隨機種子及參數數量下比較。語料為12.7M個清理後的WikiText-103詞元:

PSSA訓練交叉熵為3.98,Transformer為4.43,差距0.45 nats,困惑度分別為53.7與83.7。Transformer用盡12.7M詞元預算才達PSSA約2M詞元時的損失。

兩曲線從未交叉或接近。

訓練損失僅表示模型擬合輸入流。兩模型檢查點在未見過的198,939詞元切片上評分,PSSA始終領先,最終檢查點損失低0.51 nats。

保留差距0.43 nats與訓練差距相當,顯示PSSA非死記硬背,而是泛化能力更佳。

在相同CPU、相同提示與採樣器下生成200詞元:

遞迴模型攜帶固定大小狀態,每新增詞元成本不隨先前長度增加。Transformer每步重讀整個上下文。

為避免誤讀數據,無參數執行會顯示主畫面列出所有命令及工作目錄中找到的檢查點與語料。

上述結果在免費託管筆記本、單一入門GPU上訓練,因會話限制以20萬詞元鏈段執行。是否在10倍或100倍參數下持續優勢、記憶庫在大規模下的作用、與現代遞迴基線比較等問題,需具備真實VRAM的GPU及數天分配時間。超出入門卡的硬體將改變可探討問題範圍。

若您有計算資源,或工作單位有,這是對本專案最高效的支持。

贊助資金僅用於計算,回報為在此及成果報告中列名。請先聯絡以協商細節。

歡迎提出問題與拉取請求。最需協助部分為核心效能、現代遞迴基線比較及超越下一詞元損失的評估。請於開PR前以cargo test --release驗證分支。

直接運行依賴ureq下載數據集及tokenizers進行字節級BPE。GPU為選用。啟用--features cuda時,密集矩陣運算透過cuBLAS及設備端權重快取執行;無CUDA時使用WebGPU適配器,軟體適配器因速度慢於CPU路徑而不採用。所有功能皆可回退至CPU實作。

兩鏈各執行64個20萬詞元鏈段,每段從前一檢查點恢復,學習率排程與優化器狀態跨段持續。

匹配優化器排程對PSSA有利,對Transformer未必最佳,部分差距可能因基線訓練不足。現正進行雙模型學習率掃描,結果將公布。

訓練速率來自不同機器,無法直接比較:PSSA在Kaggle T4約900詞元/秒,基線因無GPU僅212詞元/秒。為有意義比較,兩模型在同一無GPU雙vCPU容器上,使用相同199,059詞元切片與種子42,PSSA達1716詞元/秒,基線415詞元/秒,硬體與工作量匹配下快4.1倍。早期Kaggle CPU測試為375對212。損失比較不受時間影響,因以詞元與更新數匹配。

損失為擬合流的訓練交叉熵,非保留集評估。保留集比較請參閱docs/COMPARISON.md。此規模下生成品質兩者皆差:PSSA生成「a barget of the Prian Academy」,基線生成「a material circulation of the United States」。

尚未測試兩實驗:語料切換後技能保留,以及移除512記憶槽對損失影響。

讀取TOTAL、WINDOW、FRESH環境變數並寫入--loss-csv,確保中斷後曲線保存。

位置參數與長短選項可混用。

--skip-tokens、--max-tokens與--resume結合可將長語料分段訓練,避免會話限制。若視窗跨越檔尾,選擇會從語料開頭繞回。每段訓練自身視窗並將優化器狀態傳給下一段。

kaggle/kaggle_continue.sh實現此流程,設定視窗大小與鏈段數,逐段遍歷語料並從前段檢查點恢復。status報告鏈中所有檢查點形狀與優化器步數。

DatasetManager接受一或多個逗號分隔來源。

本地檔案與目錄直接讀取;HTTP(S) URL及明確hf:owner/dataset來源會下載。結構化回應會根據text、content、article、story、instruction、output、sentence及summary等欄位合併;無支援文字欄位的結構化回應會被拒絕。

字節級BPE保留完整UTF-8大小寫、空白、標點及換行,並有完整256字元備用字母表,確保有效UTF-8不會變成<unk>。先前的小寫字詞分割器含1萬詞限制及<unk>行為,僅於--tokenizer word時可用。

下載Hugging Face數據集至本地文字檔。

網路下載不受Oxide AI驗證或策劃,訓練前請審核授權、隱私與內容。

清理提取的wikitext-103-raw文字以供新訓練使用。

同一命令可於Kaggle使用,接受本地UTF-8文字檔,非Parquet格式。-o與--out為別名。輸出路徑必須不存在(含輸入路徑或其連結),以保護原始語料。讀取、UTF-8解析及寫入失敗會透過CLI錯誤路徑非零退出,並盡可能移除部分輸出。

oxide_ai_pssa::dataset::clean_wikitext(reader, writer)為可重用庫API(BufRead/Write,回傳std::io::Result<()>)。CLI使用緩衝檔案I/O,清理器僅保留輸入/輸出行緩衝,記憶體用量與最長行長度成正比,與語料大小無關。庫呼叫者使用緩衝寫入器需自行flush;CLI明確檢查flush。無需新增依賴。

清理為選用:現有載入器、分詞器、訓練命令及kaggle/kaggle_continue.sh不變。請勿在進行中鏈段切換至清理語料,因清理改變詞元ID/數量及--skip-tokens偏移意義。請準備並一致使用同一清理語料展開新鏈段。

train命令執行兩階段訓練。

預設為潛在維度256、遞迴狀態16、記憶鍵32、記憶容量512、區塊長64、學習率1e-3、每更新8區塊、種子42。生成的二進位檔包含權重、配置、記憶、適配器與優化器狀態。非其他ML框架交換格式,應透過PSSALayer::import_from_pssa_bytes載入。

新存檔使用V7格式:包含完整V6訓練/恢復負載及有界長度前綴標準分詞器JSON。V7 BPE檢查點自包含,能還原精確排序詞彙表,無需訓練或評估語料。generate與chat拒絕V7 BPE的--data,因重新訓練分詞器無法驗證來源。V7 word與V6檢查點保留舊有可選--data精確詞彙比較。已驗證V5工件僅限推論,需--data因未含分詞器來源。

生成為自回歸,使用溫度0.70、前24候選限制後top-p 0.85過濾、最近64詞元1.25重複懲罰、立即自轉移抑制、<unk>抑制及預設64新詞元上限,遇兩句點提前結束。

V7 BPE推論還原精確嵌入分詞器,絕不從選定數據集重建。評估僅以保留文字供還原分詞器使用。

套件測試合成流中矛盾事實、MQAR式干擾、爆發重複、模型序列化及短生成提示。列印里程碑結果,未接入Cargo測試框架,非通用語言任務品質評估。

整合測試位於tests/:allocations.rs、bpe_repair.rs、checkpoint_repair.rs、core_repair.rs、linalg.rs及runtime_repair.rs,並共享工件於tests/fixtures/。涵蓋分詞器往返、跨版本檢查點匯入匯出、線性代數核心、分配行為及CLI執行輸出。Clippy無錯誤;數值核心中部分風格警告故意保留,因重寫索引迴圈會影響梯度測試穩定性。

詳見LICENSE了解專案授權。

這是一個以Rust開發的自訂AI架構。

PSSA:用Rust開發的非Transformer小型語言模型架構PSSA:用Rust開發的非Transformer小型語言模型架構PSSA:用Rust開發的非Transformer小型語言模型架構PSSA:用Rust開發的非Transformer小型語言模型架構PSSA:用Rust開發的非Transformer小型語言模型架構PSSA:用Rust開發的非Transformer小型語言模型架構