Distributed Representer in Rust: A Step-by-Step Guide

Asher Bond

Abstract

This note walks through a small Rust program that (1) encodes text into vectors, (2) applies a self-attention transform [1], and (3) stores and reloads the resulting distributed representations [2]. The code is presented as an Input–Process–Output pipeline of composable stages [3, 4, 5]. It makes the shape of a representation pipeline concrete in idiomatic Rust: three atomic transforms that compose into one run, each swappable for a production part. The encoder is bag-of-words and the attention projections are untrained, so the attention output is [ILLUSTRATIVE] structure — the pipeline's form is the subject, and every stage names the production component it stands in for.

1. Introduction

Distributed representations — dense vectors in which meaning is spread across many dimensions — are the substrate of modern retrieval and attention [1, 2]. The clearest way to understand a representation pipeline is to build a minimal one. This note does that in Rust, staged as three composable transforms: encode → attend → store/reload. The encoder is bag-of-words and the self-attention weights are untrained, so the attention output is illustrative structure that carries the pipeline's shape, not learned meaning. The program runs in one process; "distributed" refers to the representation — meaning spread across a vector's dimensions — the property this pipeline builds and persists.

2. Related Work

Self-attention. The attention block below is a stripped-down scaled dot-product attention [1]; production systems use the IO-aware FlashAttention kernel [6], not the naive softmax-over-full-matrix used here for clarity.

Vector stores. Persisting and looking up representations at scale is what vector indexes such as FAISS do [2]; the HashMap-backed store below is a teaching stand-in, not an index.

Composable stages. Framing the program as composed IPO transforms follows the functional tradition [3, 4, 5].

3. The Pipeline (Input–Process–Output)

The program is three atomic stages composed into one pipeline: store ∘ attend ∘ encode.

Stage 1 — Encode. Input: raw text. Process: build a vocabulary and count word occurrences. Output: a bag-of-words vector.

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()
 }
}

The code example here illustrates the encode stage. We define a text processor that holds a vocabulary, a simple map from each word to a position. Build vocab walks over the training texts and gives every new word its own index. Then text to vector counts how often each known word appears and returns those counts as a single row. In short, it is a bag of words: order is dropped, but the presence and frequency of each term is kept.

Stage 2 — Attend. Input: a vector. Process: scaled dot-product self-attention with random (untrained) Q/K/V projections [1]. Output: a context vector. This is illustrative: with random weights the output carries no learned meaning, and the naive softmax-over-full-scores is the exact cost FlashAttention [6] exists to reduce.

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)
 }
}

Stage 3 — Store / reload. Input: a key and a representation. Process: insert into an in-memory map; serialize to / from JSON. Output: a persisted, reloadable representation. (A production system uses a vector index [2], not a HashMap.)

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()
 }
}

Composition — main.rs. The three stages compose into one run:

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."),
 }
}

Run with cargo run. The listing compiles and runs as written: save imports BufWriter and constructs the writer by value, and the softmax reuses a single exp array across the normalization.

Evidence & Scope

This is a reference implementation with a precise remit: make the Input–Process–Output shape of a representation pipeline concrete in idiomatic Rust [4], with each stage swappable for the production part it stands in for. The program is single-process — "distributed" describes the vector representation, not the execution. The attention projections are untrained and the encoder is bag-of-words, so the attention output is [ILLUSTRATIVE] structure that carries the pipeline's form rather than learned meaning. The HashMap store does exact key lookup; a real index [2] adds the ANN properties that turn lookup into similarity search. The softmax-over-full-scores block is the exact cost FlashAttention [6] removes. Each of these is the seam where a production component drops in — the pipeline is built to be extended along them.

References

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.