optirs_nas/multimodal_nas/
types.rs

1//! Core value types for the multimodal NAS module: modalities, layer and
2//! fusion-operator enums, the per-modality encoder, and the helper
3//! functions that classify layers and fusion operators.
4
5use serde::{Deserialize, Serialize};
6use std::fmt;
7
8// =======================================================================
9// Modalities
10// =======================================================================
11
12/// Input modality consumed by a multimodal architecture.
13///
14/// `Modality` is `Ord` so that declared-modality lists have a stable,
15/// deterministic ordering and so the type can key ordered collections.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
17pub enum Modality {
18    /// Visual input (pixels / feature maps).
19    Image,
20    /// Textual / token input.
21    Text,
22    /// Audio waveform / spectrogram input.
23    Audio,
24}
25
26impl fmt::Display for Modality {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        match self {
29            Modality::Image => write!(f, "image"),
30            Modality::Text => write!(f, "text"),
31            Modality::Audio => write!(f, "audio"),
32        }
33    }
34}
35
36// =======================================================================
37// Activation kinds
38// =======================================================================
39
40/// Non-linear activation applied by [`MultimodalLayer::Activation`].
41///
42/// Kept as a fieldless enum so [`MultimodalLayer`] can derive `Eq` / `Hash`.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
44pub enum ActivationKind {
45    /// Rectified linear unit.
46    Relu,
47    /// Gaussian error linear unit.
48    Gelu,
49    /// Hyperbolic tangent.
50    Tanh,
51    /// Logistic sigmoid.
52    Sigmoid,
53}
54
55// =======================================================================
56// Layers
57// =======================================================================
58
59/// Encoder / head layer types available in multimodal architectures.
60///
61/// Each variant carries the hyperparameters that influence its compute /
62/// memory footprint. Concrete fields use `u32` (rather than `usize`)
63/// because configurations are persisted to JSON and compared for equality
64/// across machines with different pointer widths.
65///
66/// Layers fall into three groups:
67///
68/// * *Image* encoders ([`MultimodalLayer::ImageConv`],
69///   [`MultimodalLayer::ImagePool`]),
70/// * *Text* encoders ([`MultimodalLayer::TextEmbed`],
71///   [`MultimodalLayer::TextTransformer`]),
72/// * *Audio* encoders ([`MultimodalLayer::AudioConv`],
73///   [`MultimodalLayer::AudioRecurrent`]),
74///
75/// plus the *generic* layers ([`MultimodalLayer::Dense`],
76/// [`MultimodalLayer::Norm`], [`MultimodalLayer::Activation`]) which may
77/// appear in any modality's encoder and in the head. The mapping from a
78/// layer to its owning modality (or `None` for generic layers) is exposed
79/// by [`layer_modality`].
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
81pub enum MultimodalLayer {
82    /// 2D convolution for the image encoder.
83    ImageConv {
84        /// Number of output channels.
85        channels: u32,
86        /// Square convolution kernel size.
87        kernel: u32,
88    },
89    /// Spatial pooling for the image encoder.
90    ImagePool {
91        /// Pooling window size.
92        kernel: u32,
93    },
94    /// Token-embedding layer for the text encoder.
95    TextEmbed {
96        /// Embedding dimension.
97        dim: u32,
98    },
99    /// Transformer block for the text encoder.
100    TextTransformer {
101        /// Number of attention heads.
102        heads: u32,
103        /// Per-head feature dimension.
104        dim: u32,
105    },
106    /// 1D convolution for the audio encoder.
107    AudioConv {
108        /// Number of output channels.
109        channels: u32,
110        /// Convolution kernel size (in frames).
111        kernel: u32,
112    },
113    /// Recurrent cell for the audio encoder.
114    AudioRecurrent {
115        /// Hidden state size.
116        hidden: u32,
117    },
118    /// Generic dense projection (usable in any encoder or in the head).
119    Dense {
120        /// Output feature dimension.
121        out: u32,
122    },
123    /// Generic normalisation layer.
124    Norm,
125    /// Generic non-linear activation.
126    Activation {
127        /// The activation function applied.
128        kind: ActivationKind,
129    },
130}
131
132/// Return the modality a layer belongs to, or `None` for generic layers.
133///
134/// This is the single source of truth for encoder membership in the
135/// multimodal-NAS module — [`super::MultimodalArchitecture::validate`] uses
136/// it to reject foreign layers placed in the wrong encoder.
137pub fn layer_modality(layer: &MultimodalLayer) -> Option<Modality> {
138    match layer {
139        MultimodalLayer::ImageConv { .. } | MultimodalLayer::ImagePool { .. } => {
140            Some(Modality::Image)
141        }
142        MultimodalLayer::TextEmbed { .. } | MultimodalLayer::TextTransformer { .. } => {
143            Some(Modality::Text)
144        }
145        MultimodalLayer::AudioConv { .. } | MultimodalLayer::AudioRecurrent { .. } => {
146            Some(Modality::Audio)
147        }
148        MultimodalLayer::Dense { .. }
149        | MultimodalLayer::Norm
150        | MultimodalLayer::Activation { .. } => None,
151    }
152}
153
154/// Rough compute / parameter "capacity" of a single layer.
155///
156/// The value is a unitless cost used solely for the modality-balance
157/// constraint; it is intentionally monotone in the dominant size
158/// hyperparameter of each layer. Every layer has capacity `>= 1.0`, so
159/// encoder capacities are always strictly positive.
160pub fn layer_capacity(layer: &MultimodalLayer) -> f64 {
161    match layer {
162        MultimodalLayer::ImageConv { channels, .. } => f64::from(*channels),
163        MultimodalLayer::ImagePool { .. } => 8.0,
164        MultimodalLayer::TextEmbed { dim } => f64::from(*dim),
165        MultimodalLayer::TextTransformer { heads, dim } => f64::from(*heads) * f64::from(*dim),
166        MultimodalLayer::AudioConv { channels, .. } => f64::from(*channels),
167        MultimodalLayer::AudioRecurrent { hidden } => f64::from(*hidden),
168        MultimodalLayer::Dense { out } => f64::from(*out),
169        MultimodalLayer::Norm => 4.0,
170        MultimodalLayer::Activation { .. } => 1.0,
171    }
172}
173
174// =======================================================================
175// Fusion operators
176// =======================================================================
177
178/// Fusion operators combining the per-modality encoder outputs.
179#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
180pub enum FusionOp {
181    /// Early fusion: concatenate encoder features before the head.
182    EarlyConcat,
183    /// Late fusion: combine per-modality predictions (e.g. averaging).
184    LateFusion,
185    /// Cross-modal attention between encoder streams.
186    CrossAttention {
187        /// Number of attention heads.
188        heads: u32,
189        /// Per-head feature dimension.
190        dim: u32,
191    },
192    /// Gated fusion: one modality gates another.
193    Gated {
194        /// Gate / projection dimension.
195        dim: u32,
196    },
197    /// Bilinear pooling of exactly two modality features.
198    Bilinear {
199        /// Bilinear projection dimension.
200        dim: u32,
201    },
202}
203
204/// Minimum number of modalities a fusion operator can combine.
205pub fn fusion_min_modalities(fusion: &FusionOp) -> usize {
206    match fusion {
207        FusionOp::EarlyConcat | FusionOp::LateFusion => 1,
208        FusionOp::CrossAttention { .. } | FusionOp::Gated { .. } | FusionOp::Bilinear { .. } => 2,
209    }
210}
211
212/// Maximum number of modalities a fusion operator can combine, if bounded.
213pub fn fusion_max_modalities(fusion: &FusionOp) -> Option<usize> {
214    match fusion {
215        FusionOp::Bilinear { .. } => Some(2),
216        _ => None,
217    }
218}
219
220/// Return `true` iff `fusion` is compatible with `num_modalities` inputs.
221pub fn fusion_supports(fusion: &FusionOp, num_modalities: usize) -> bool {
222    let min_m = fusion_min_modalities(fusion);
223    match fusion_max_modalities(fusion) {
224        None => num_modalities >= min_m,
225        Some(max_m) => num_modalities >= min_m && num_modalities <= max_m,
226    }
227}
228
229// =======================================================================
230// Modality encoder
231// =======================================================================
232
233/// A per-modality encoder: the modality it encodes and its layer sequence.
234#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
235pub struct ModalityEncoder {
236    /// The modality this encoder consumes.
237    pub modality: Modality,
238    /// Layers applied in execution order.
239    pub layers: Vec<MultimodalLayer>,
240}
241
242/// Sum of the per-layer capacities of an encoder.
243pub fn encoder_capacity(encoder: &ModalityEncoder) -> f64 {
244    encoder.layers.iter().map(layer_capacity).sum()
245}
optirs_nas/multimodal_nas/types.rs

optirs_nas/multimodal_nas/
types.rs
