voxtral_micro/models/
time_embedding.rs

1//! Time embedding for Voxtral Realtime.
2//!
3//! Sinusoidal embedding that encodes the transcription delay.
4
5use burn::prelude::*;
6
7/// Time embedding module that produces sinusoidal embeddings.
8///
9/// Used to encode the transcription delay as a conditioning signal
10/// for the ADA RMSNorm modulation in the decoder layers.
11#[derive(Debug)]
12pub struct TimeEmbedding {
13    /// Dimension of the embedding
14    dim: usize,
15    /// Base frequency (default: 10000.0)
16    theta: f32,
17}
18
19impl TimeEmbedding {
20    /// Create a new time embedding with given dimension.
21    pub fn new(dim: usize) -> Self {
22        Self {
23            dim,
24            theta: 10000.0,
25        }
26    }
27
28    /// Create a new time embedding with custom theta.
29    pub fn with_theta(dim: usize, theta: f32) -> Self {
30        Self { dim, theta }
31    }
32
33    /// Compute sinusoidal embedding for a time value.
34    ///
35    /// # Arguments
36    /// * `t` - Time value (typically the number of delay tokens)
37    /// * `device` - Device to create the tensor on
38    ///
39    /// # Returns
40    /// Tensor of shape [1, 1, dim] containing the sinusoidal embedding
41    pub fn embed<B: Backend>(&self, t: f32, device: &B::Device) -> Tensor<B, 3> {
42        let half_dim = self.dim / 2;
43
44        // Compute inverse frequencies: exp(-log(theta) * i / (dim/2)) for i in 0..dim/2
45        let mut inv_freq = Vec::with_capacity(half_dim);
46        let log_theta = self.theta.ln();
47        for i in 0..half_dim {
48            let freq = (-log_theta * (i as f32) / (half_dim as f32)).exp();
49            inv_freq.push(freq);
50        }
51
52        // Compute t * inv_freq
53        let mut cos_vals = Vec::with_capacity(half_dim);
54        let mut sin_vals = Vec::with_capacity(half_dim);
55        for &freq in &inv_freq {
56            let angle = t * freq;
57            cos_vals.push(angle.cos());
58            sin_vals.push(angle.sin());
59        }
60
61        // Concatenate [cos, sin] to get full embedding
62        let mut embedding = Vec::with_capacity(self.dim);
63        embedding.extend_from_slice(&cos_vals);
64        embedding.extend_from_slice(&sin_vals);
65
66        // Create tensor with shape [1, 1, dim]
67        Tensor::from_data(
68            burn::tensor::TensorData::new(embedding, [1, 1, self.dim]),
69            device,
70        )
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77    use burn::backend::Wgpu;
78
79    type TestBackend = Wgpu;
80
81    #[test]
82    fn test_time_embedding_shape() {
83        let device = Default::default();
84        let embed = TimeEmbedding::new(3072);
85
86        let t_cond = embed.embed::<TestBackend>(6.0, &device);
87        assert_eq!(t_cond.dims(), [1, 1, 3072]);
88    }
89
90    #[test]
91    fn test_time_embedding_values() {
92        let device = Default::default();
93        let embed = TimeEmbedding::new(4);
94
95        let t_cond = embed.embed::<TestBackend>(1.0, &device);
96        let data = t_cond.to_data();
97        let slice = data.as_slice::<f32>().unwrap();
98
99        // For t=1, dim=4:
100        // inv_freq[0] = exp(-log(10000) * 0 / 2) = exp(0) = 1.0
101        // inv_freq[1] = exp(-log(10000) * 1 / 2) = exp(-log(10000)/2) = 1/100 = 0.01
102        // cos(1 * 1.0) = cos(1) ≈ 0.5403
103        // cos(1 * 0.01) = cos(0.01) ≈ 0.99995
104        // sin(1 * 1.0) = sin(1) ≈ 0.8415
105        // sin(1 * 0.01) = sin(0.01) ≈ 0.01
106
107        // Check approximate values
108        assert!(
109            (slice[0] - 0.5403).abs() < 0.01,
110            "cos(1) wrong: {}",
111            slice[0]
112        );
113        assert!(
114            (slice[1] - 0.99995).abs() < 0.001,
115            "cos(0.01) wrong: {}",
116            slice[1]
117        );
118        assert!(
119            (slice[2] - 0.8415).abs() < 0.01,
120            "sin(1) wrong: {}",
121            slice[2]
122        );
123        assert!(
124            (slice[3] - 0.01).abs() < 0.001,
125            "sin(0.01) wrong: {}",
126            slice[3]
127        );
128    }
129
130    #[test]
131    fn test_time_embedding_vs_python() {
132        // Test against expected Python output
133        // Python: TimeEmbedding(dim=8, theta=10000)(torch.tensor([6.0]))
134        // Should produce cos and sin of 6 * inv_freq
135        let device = Default::default();
136        let embed = TimeEmbedding::new(8);
137
138        let t_cond = embed.embed::<TestBackend>(6.0, &device);
139        let data = t_cond.to_data();
140        let slice = data.as_slice::<f32>().unwrap();
141
142        // inv_freq = [1.0, 0.1, 0.01, 0.001] (approx for theta=10000, dim=8)
143        // 6 * inv_freq = [6.0, 0.6, 0.06, 0.006]
144        // cos(6) ≈ 0.9602
145        // cos(0.6) ≈ 0.8253
146        // cos(0.06) ≈ 0.9982
147        // cos(0.006) ≈ 0.99998
148
149        println!("t_cond for t=6: {:?}", slice);
150        assert!(slice.len() == 8);
151    }
152}
voxtral_micro/models/time_embedding.rs

voxtral_micro/models/
time_embedding.rs
