uqa_ml/
deep_fusion.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Multi-layer fusion operator (Section 7, Paper 4).
8//!
9//! Implements deep Bayesian fusion as a multi-layer network:
10//!
11//! ```text
12//!     l^(k) = g( l^(k-1) + sum_j logit(P_j^(k)) )
13//!     P_final = sigmoid(l^(K))
14//! ```
15//!
16//! The internal channel map keys per-document feature vectors. Each
17//! `Layer` variant updates that map, including signal, dense, convolutional,
18//! recurrent, normalization, attention, and graph-aware propagation / pooling
19//! layers.
20
21use std::collections::BTreeMap;
22use std::sync::Arc;
23
24use uqa_core::{IndexStats, Payload, PostingEntry, PostingList, Value};
25use uqa_scoring::prob::{confidence_scaled_log_odds_pool_weighted, logit, sigmoid, PROB_EPSILON};
26
27use uqa_operators::{
28    base::{Direction, OperatorResult},
29    ExecutionContext, Operator,
30};
31use uqa_storage::{StorageBackendError, StorageBackendResult};
32
33use crate::backend::{try_filled_vec, try_vec_with_capacity, MLError, MLResult};
34
35#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
36pub enum Gating {
37    #[default]
38    None,
39    Softplus,
40    Sigmoid,
41    ReLU,
42    Swish,
43    Gelu,
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum GlobalPoolMethod {
48    Avg,
49    Max,
50    AvgMax,
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub enum AggregationKind {
55    Mean,
56    Sum,
57    Max,
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub enum PoolMethod {
62    Avg,
63    Max,
64}
65
66#[allow(clippy::upper_case_acronyms)]
67#[derive(Clone)]
68pub enum Layer {
69    /// Runtime-provided feature vector. This layer is a no-op during
70    /// execution; it marks the expected input dimension for trained
71    /// models that receive feature batches from an ML backend.
72    Input { dimensions: usize },
73    /// Run a list of `Operator` signals, fuse them via log-odds
74    /// conjunction at the configured `alpha`, then add the resulting
75    /// logit to channel 0 as a residual connection.
76    Signal(Vec<Arc<dyn Operator>>),
77    /// Initialize the channel map from a raw embedding vector. Element
78    /// `i` becomes node `i+1` with a single-channel value.
79    Embed(Vec<f64>),
80    /// Fully connected: `out = W @ input + bias`, then gating.
81    Dense {
82        /// `output_channels x input_channels`, row-major.
83        weights: Vec<f64>,
84        bias: Vec<f64>,
85        output_channels: usize,
86        input_channels: usize,
87    },
88    /// Concatenate every node's channel vector into a single vector.
89    Flatten,
90    /// Reduce all spatial nodes to one vector.
91    GlobalPool(GlobalPoolMethod),
92    /// Numerically stable softmax per node.
93    Softmax,
94    /// Per-channel batch normalization across all nodes.
95    BatchNorm { epsilon: f64 },
96    /// Inference-mode dropout: scale every value by `1 - p`.
97    Dropout { p: f64 },
98    /// One-dimensional CNN over sorted sequence positions.
99    ///
100    /// Weights are row-major as `output_channels x kernel_size x input_channels`.
101    CNN1D {
102        weights: Vec<f64>,
103        bias: Vec<f64>,
104        output_channels: usize,
105        input_channels: usize,
106        kernel_size: usize,
107        stride: usize,
108        padding: usize,
109    },
110    /// Two-dimensional CNN over flattened `H x W x C` spatial positions.
111    ///
112    /// Weights are row-major as
113    /// `output_channels x kernel_height x kernel_width x input_channels`.
114    CNN2D {
115        weights: Vec<f64>,
116        bias: Vec<f64>,
117        output_channels: usize,
118        input_channels: usize,
119        input_height: usize,
120        input_width: usize,
121        kernel_height: usize,
122        kernel_width: usize,
123        stride_height: usize,
124        stride_width: usize,
125        padding_height: usize,
126        padding_width: usize,
127    },
128    /// Propagate channel-0 scores through graph edges.
129    ///
130    /// `aggregation` averages / sums / maxes the in-bounds neighbor
131    /// probabilities; the resulting logit is added as a residual on
132    /// channel 0. Requires `ExecutionContext::graph`.
133    Propagate {
134        /// Edge label to follow. An empty string selects every edge label.
135        edge_label: String,
136        aggregation: AggregationKind,
137        direction: Direction,
138    },
139    /// Weighted multi-hop graph convolution on channel 0.
140    ///
141    /// `hop_weights[0]` is the self weight, `hop_weights[i]` weights
142    /// the average over the hop-`i` neighbor ring. Weights are
143    /// L1-normalized; the result is converted back to logit and added
144    /// as a residual.
145    Conv {
146        /// Edge label to follow. An empty string selects every edge label.
147        edge_label: String,
148        hop_weights: Vec<f64>,
149        direction: Direction,
150    },
151    /// Spatial downsampling via greedy BFS partitioning.
152    ///
153    /// Groups `pool_size` neighboring nodes via BFS, aggregates their
154    /// channel vectors element-wise (`PoolMethod::{Avg, Max}`), and
155    /// keeps the smallest doc id as the representative.
156    Pool {
157        /// Edge label to follow. An empty string selects every edge label.
158        edge_label: String,
159        pool_size: usize,
160        method: PoolMethod,
161        direction: Direction,
162    },
163    /// Self-attention across the per-node channel vectors with
164    /// `Q = K = V = X`, scaled-dot-product, no learned projections.
165    Attention,
166    /// Vanilla RNN over sorted sequence positions.
167    ///
168    /// Weights are row-major as `hidden_channels x input_channels` and
169    /// `hidden_channels x hidden_channels`.
170    RNN {
171        weights_input: Vec<f64>,
172        weights_hidden: Vec<f64>,
173        bias: Vec<f64>,
174        hidden_channels: usize,
175        input_channels: usize,
176        return_sequences: bool,
177    },
178    /// LSTM over sorted sequence positions.
179    ///
180    /// Gate order is input, forget, candidate, output. Both weight
181    /// matrices are row-major with `4 * hidden_channels` rows.
182    LSTM {
183        weights_input: Vec<f64>,
184        weights_hidden: Vec<f64>,
185        bias: Vec<f64>,
186        hidden_channels: usize,
187        input_channels: usize,
188        return_sequences: bool,
189    },
190}
191
192pub struct DeepFusionOperator {
193    layers: Vec<Layer>,
194    alpha: f64,
195    gating: Gating,
196}
197
198impl DeepFusionOperator {
199    pub fn new(layers: Vec<Layer>, alpha: f64, gating: Gating) -> MLResult<Self> {
200        validate_layers(&layers, alpha)?;
201        Ok(Self {
202            layers,
203            alpha,
204            gating,
205        })
206    }
207
208    pub fn layers(&self) -> &[Layer] {
209        &self.layers
210    }
211
212    pub fn alpha(&self) -> f64 {
213        self.alpha
214    }
215
216    pub fn gating(&self) -> Gating {
217        self.gating
218    }
219}
220
221mod attention;
222mod cnn;
223mod execution;
224mod graph_layers;
225mod layer_dispatch;
226mod recurrent;
227mod runtime;
228mod state;
229mod tensor_layers;
230mod validation;
231
232use attention::{apply_attention, build_result};
233use cnn::{apply_cnn_1d, apply_cnn_2d};
234use graph_layers::{apply_conv, apply_pool, apply_propagate};
235use layer_dispatch::{apply_cnn_1d_layer, apply_cnn_2d_layer, apply_lstm_layer, apply_rnn_layer};
236use recurrent::{apply_lstm, apply_rnn};
237use runtime::{
238    apply_gating, runtime_filled_vec, runtime_model_error, runtime_vec_with_capacity, safe_logit,
239    usize_to_f64_exact,
240};
241use state::{Convolution1D, Convolution2D, ForwardState, LongShortTermMemory, Recurrent};
242use tensor_layers::{
243    apply_batch_norm, apply_dense, apply_dropout, apply_embed, apply_flatten, apply_global_pool,
244    apply_signal, apply_softmax,
245};
246use validation::{validate_layers, validate_state};
247
248#[cfg(test)]
249mod tests;
uqa_ml/deep_fusion.rs

uqa_ml/
deep_fusion.rs
