Rustで作る分散表現器: ステップバイステップ実装

Asher Bond

Rust による最小の分散表現ストア——単一ノードの参照実装

著者: Asher Bond (asher.bond@distillative.ai)

要旨

本稿は、(1) テキストをベクトルへ符号化し、(2) 自己注意変換 [1] を適用し、(3) 得られた分散表現を保存・再読込する、小さな Rust プログラムを一通り追う [2]。コードは、合成可能な段を連ねた Input–Process–Output パイプラインとして提示する [3, 4, 5]。これにより、表現パイプラインの「かたち」を、慣用的な Rust で具体化する。三つの原子的変換が一つの実行へと合成され、各段はそれぞれ本番用の部品へ差し替え可能である。符号化器は bag-of-words であり、注意の射影は未学習であるから、注意出力は [ILLUSTRATIVE] な構造にすぎない——主題はパイプラインの形式であって、各段はそれが代替する本番コンポーネントを名指しする。

1. 序論

分散表現——意味が多数の次元にわたって分散した密ベクトル——は、現代の検索と注意の基盤である [1, 2]。表現パイプラインを理解する最も明快な方法は、最小のものを一つ組み上げることである。本稿はそれを Rust で行い、三つの合成可能な変換として段階化する。すなわち 符号化 → 注意 → 保存/再読込 である。符号化器は bag-of-words、自己注意の重みは未学習であるから、注意出力は学習された意味ではなくパイプラインの形状を担う説明用の構造である。プログラムは単一プロセスで動く。「分散」が指すのは「表現」——ベクトルの各次元に分散した意味——であり、このパイプラインが構築し永続化する当の性質である。

2. 関連研究

自己注意。 以下の注意ブロックは、スケール化ドット積注意 [1] を切り詰めたものである。本番システムでは、ここで明快さのために用いる素朴な「全行列にわたる softmax」ではなく、IO を意識した FlashAttention カーネル [6] を用いる。

ベクトルストア。 表現を大規模に永続化し検索するのが、FAISS のようなベクトルインデックスの役割である [2]。以下の HashMap を裏づけとするストアは教育用の代役であって、インデックスではない。

合成可能な段。 プログラムを合成された IPO 変換として枠づける発想は、関数型の伝統に連なる [3, 4, 5]。

3. パイプライン (Input–Process–Output)

プログラムは、一つのパイプラインへ合成された三つの原子的な段である。すなわち store ∘ attend ∘ encode。

第1段——符号化。 入力: 生テキスト。処理: 語彙を構築し、単語の出現を数える。出力: bag-of-words ベクトル。

use ndarray::Array2;
use std::collections::HashMap;

pub struct TextProcessor {
 pub vocab: HashMap<String, usize>,
}

impl TextProcessor {
 pub fn new() -> Self {
 TextProcessor { vocab: HashMap::new() }
 }

 pub fn build_vocab(&mut self, texts: &[&str]) {
 let mut index = 0;
 for &text in texts {
 for word in text.split_whitespace() {
 if !self.vocab.contains_key(word) {
 self.vocab.insert(word.to_string(), index);
 index += 1;
 }
 }
 }
 }

 pub fn text_to_vector(&self, text: &str) -> Array2<f32> {
 let mut vector = vec![0.0; self.vocab.len()];
 for word in text.split_whitespace() {
 if let Some(&index) = self.vocab.get(word) {
 vector[index] += 1.0;
 }
 }
 Array2::from_shape_vec((1, self.vocab.len()), vector).unwrap()
 }
}

第2段——注意。 入力: ベクトル。処理: ランダム（未学習）な Q/K/V 射影によるスケール化ドット積自己注意 [1]。出力: 文脈ベクトル。これは説明用である。ランダムな重みのもとで出力は学習された意味を担わず、素朴な「全スコアにわたる softmax」こそ FlashAttention [6] が削減するために存在する、まさにそのコストである。

use ndarray::{Array2, Axis};
use ndarray_rand::RandomExt;
use rand::distributions::Uniform;

pub struct SelfAttention {
 pub query: Array2<f32>,
 pub key: Array2<f32>,
 pub value: Array2<f32>,
 pub output_weight: Array2<f32>,
}

impl SelfAttention {
 pub fn new(dim: usize) -> Self {
 let distribution = Uniform::new(-0.1, 0.1);
 SelfAttention {
 query: Array2::random((dim, dim), distribution),
 key: Array2::random((dim, dim), distribution),
 value: Array2::random((dim, dim), distribution),
 output_weight: Array2::random((dim, dim), distribution),
 }
 }

 pub fn attention(&self, x: &Array2<f32>) -> Array2<f32> {
 let q = x.dot(&self.query);
 let k = x.dot(&self.key);
 let v = x.dot(&self.value);

 let scores = q.dot(&k.t()) / (x.shape()[1] as f32).sqrt();
 let exp = scores.mapv(|s| s.exp());
 let weights = &exp / &exp.sum_axis(Axis(1)).insert_axis(Axis(1));
 let context = weights.dot(&v);
 context.dot(&self.output_weight)
 }
}

第3段——保存/再読込。 入力: キーと表現。処理: インメモリのマップへ挿入し、JSON へ／から直列化する。出力: 永続化され再読込可能な表現。（本番システムでは HashMap ではなくベクトルインデックス [2] を用いる。）

use ndarray::Array2;
use std::collections::HashMap;
use serde::{Serialize, Deserialize};
use std::fs::File;
use std::io::{BufWriter, Read};

#[derive(Serialize, Deserialize)]
pub struct DistributedRepresenter {
 pub representations: HashMap<String, Array2<f32>>,
}

impl DistributedRepresenter {
 pub fn new() -> Self {
 DistributedRepresenter { representations: HashMap::new() }
 }

 pub fn add_representation(&mut self, key: String, representation: Array2<f32>) {
 self.representations.insert(key, representation);
 }

 pub fn get_representation(&self, key: &str) -> Option<&Array2<f32>> {
 self.representations.get(key)
 }

 pub fn save(&self, path: &str) {
 let file = File::create(path).unwrap();
 let writer = BufWriter::new(file);
 serde_json::to_writer(writer, &self).unwrap();
 }

 pub fn load(path: &str) -> Self {
 let mut file = File::open(path).unwrap();
 let mut data = String::new();
 file.read_to_string(&mut data).unwrap();
 serde_json::from_str(&data).unwrap()
 }
}

合成——main.rs。 三つの段が一つの実行へと合成される。

mod attention;
mod distributed_representation;
mod input_processing;

use attention::SelfAttention;
use distributed_representation::DistributedRepresenter;
use input_processing::TextProcessor;

fn main() {
 let mut text_processor = TextProcessor::new();
 let texts = vec!["hello world", "machine learning is fun", "rust is great for performance"];
 text_processor.build_vocab(&texts);

 let input_vector = text_processor.text_to_vector("machine learning in rust");
 println!("Input Vector:\n{:?}", input_vector);

 let dim = text_processor.vocab.len();
 let self_attention = SelfAttention::new(dim);
 let attention_output = self_attention.attention(&input_vector);
 println!("Attention Output:\n{:?}", attention_output);

 let mut representer = DistributedRepresenter::new();
 representer.add_representation("example".to_string(), attention_output.clone());
 representer.save("representations.json");

 let loaded = DistributedRepresenter::load("representations.json");
 match loaded.get_representation("example") {
 Some(r) => println!("Retrieved Representation:\n{:?}", r),
 None => println!("Representation not found."),
 }
}

cargo run で実行する。このリストは書かれたままでコンパイルされ動作する。save は BufWriter を取り込んで書き手を値で構築し、softmax は正規化を通じて単一の exp 配列を使い回す。

根拠と適用範囲

これは、明確な使命をもつ参照実装である。すなわち、表現パイプラインの Input–Process–Output という形状を慣用的な Rust [4] で具体化し、各段を、それが代替する本番用の部品へ差し替え可能にすることである。プログラムは単一プロセスであり、「分散」が言い表すのは実行ではなくベクトル表現である。注意の射影は未学習であり符号化器は bag-of-words であるから、注意出力は学習された意味ではなくパイプラインの形式を担う [ILLUSTRATIVE] な構造である。HashMap ストアは厳密なキー照合を行う。実際のインデックス [2] は、照合を類似検索へと変える ANN 的性質を加える。「全スコアにわたる softmax」のブロックは、FlashAttention [6] が取り除く、まさにそのコストである。これらの一つひとつが、本番コンポーネントが差し込まれる継ぎ目であり——パイプラインはそれらに沿って拡張されるべく組まれている。

参考文献

Ashish Vaswani et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). arXiv:1706.03762.
Jeff Johnson et al. (2019). Billion-Scale Similarity Search with GPUs. IEEE Transactions on Big Data. arXiv:1702.08734. [FAISS]
John Backus (1978). Can Programming Be Liberated from the von Neumann Style? A Functional Style and Its Algebra of Programs. Communications of the ACM. [1977 ACM Turing Award Lecture]
John Hughes (1989). Why Functional Programming Matters. The Computer Journal.
Christopher Strachey (2000). Fundamental Concepts in Programming Languages. Higher-Order and Symbolic Computation. [Reprint of 1967 lecture notes]
Tri Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems (NeurIPS). arXiv:2205.14135.