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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
// Quantum Token Embedding — maps tokens to density matrix states.
//
// Each token is embedded as a parameterized pure state on the Bloch hypersphere.
// Similar tokens → nearby quantum states (high fidelity).
// Dissimilar tokens → orthogonal states (zero fidelity).
use crate::complex::Complex;
use crate::density_matrix::DensityMatrixN;
/// Quantum embedding: maps vocab indices to density matrices.
#[derive(Clone)]
pub struct QuantumEmbedding {
pub vocab_size: usize,
pub dim: usize,
/// Per-token rotation angles [vocab_size × dim]. Learnable.
/// Each token gets `dim` angles that parameterize a unitary rotation.
pub angles: Vec<f32>,
}
impl QuantumEmbedding {
pub fn new(vocab_size: usize, dim: usize, seed: u64) -> Self {
let mut angles = Vec::with_capacity(vocab_size * dim);
for i in 0..(vocab_size * dim) {
let s = seed.wrapping_add(i as u64).wrapping_mul(0x517cc1b727220a95);
angles.push(std::f32::consts::PI * ((s % 2000) as f32 / 1000.0 - 1.0));
}
Self {
vocab_size,
dim,
angles,
}
}
/// Embed a single token as a density matrix.
/// Start with |k mod dim⟩, then apply parameterized rotations to spread
/// the state across modes. The rotations are the learnable embedding.
pub fn embed(&self, token: usize) -> DensityMatrixN {
let d = self.dim;
let base = token % d;
// Build amplitudes from angles via parameterized rotation
let offset = token * d;
let mut amplitudes = vec![Complex::ZERO; d];
// Start with base state
amplitudes[base] = Complex::ONE;
// Apply rotation: each angle mixes the base with another mode
for k in 0..d {
if k == base {
continue;
}
let angle = self.angles.get(offset + k).copied().unwrap_or(0.0);
let c = angle.cos();
let s = angle.sin();
let old_base = amplitudes[base];
let old_k = amplitudes[k];
amplitudes[base] = Complex::new(c * old_base.re - s * old_k.re, c * old_base.im - s * old_k.im);
amplitudes[k] = Complex::new(s * old_base.re + c * old_k.re, s * old_base.im + c * old_k.im);
}
// Build density matrix ρ = |ψ⟩⟨ψ|
let mut entries = vec![Complex::ZERO; d * d];
for i in 0..d {
for j in 0..d {
entries[i * d + j] = amplitudes[i].mul(amplitudes[j].conj());
}
}
let n2 = d * d;
DensityMatrixN {
dim: d,
entries,
scratch_a: vec![Complex::ZERO; n2],
scratch_b: vec![Complex::ZERO; n2],
}
}
/// Embed a single token as an amplitude vector (Fock space representation).
///
/// Returns the dim-dimensional complex state vector |ψ⟩ WITHOUT forming
/// the density matrix ρ = |ψ⟩⟨ψ|. This is the Fock space representation:
/// dim amplitudes instead of dim² matrix entries.
///
/// BA-62: At dim=86, this returns 688 bytes instead of 59,168 bytes (86× smaller).
/// The populations |ψ_k|² are identical to the diagonal of ρ = |ψ⟩⟨ψ|.
pub fn embed_amplitude(&self, token: usize) -> Vec<Complex> {
let d = self.dim;
let base = token % d;
let offset = token * d;
let mut amplitudes = vec![Complex::ZERO; d];
amplitudes[base] = Complex::ONE;
for k in 0..d {
if k == base {
continue;
}
let angle = self.angles.get(offset + k).copied().unwrap_or(0.0);
let c = angle.cos();
let s = angle.sin();
let old_base = amplitudes[base];
let old_k = amplitudes[k];
amplitudes[base] = Complex::new(c * old_base.re - s * old_k.re, c * old_base.im - s * old_k.im);
amplitudes[k] = Complex::new(s * old_base.re + c * old_k.re, s * old_base.im + c * old_k.im);
}
amplitudes
}
/// Number of learnable parameters.
pub fn num_params(&self) -> usize {
self.angles.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn embedding_produces_valid_state() {
let emb = QuantumEmbedding::new(65, 5, 42);
for token in 0..65 {
let rho = emb.embed(token);
assert!(
(rho.trace() - 1.0).abs() < 1e-4,
"token {token}: trace = {}",
rho.trace()
);
assert!(
rho.purity() > 0.9,
"token {token}: embedded state should be nearly pure, got {}",
rho.purity()
);
}
}
#[test]
fn different_tokens_different_states() {
let emb = QuantumEmbedding::new(65, 5, 42);
let rho_a = emb.embed(0);
let rho_b = emb.embed(1);
// Check that populations differ
let pops_a = rho_a.populations();
let pops_b = rho_b.populations();
let diff: f32 = pops_a.iter().zip(pops_b.iter()).map(|(a, b)| (a - b).abs()).sum();
assert!(diff > 0.01, "Different tokens should produce different states");
}
}