mirror of
https://github.com/ruvnet/RuView
synced 2026-08-11 20:41:44 +00:00
fix(train,nn): Tier-2 correctness/security — metric scale, OOM bounds, panics (ADR-155 §Tier-2)
Each fix ships a test that would have caught the bug: - ruview_metrics OKS: derive scale from GT extent (no s=1.0 fake-Gold), reject s<=0, bound the loop to array extents (no panic on short/adversarial input). - config.validate(): UPPER bounds on window_frames/subcarriers/backbone_channels/ heatmap_size/keypoints/body_parts/batch_size + reject negative gpu_device_id (closes the config-OOM class); defaults+presets still validate. - subcarrier.rs: graceful fallback instead of panic on non-contiguous input. - ablation.rs latency_percentiles: total_cmp + NaN guard (no partial_cmp unwrap). - tensor.rs softmax(axis): normalize per-lane along the given axis (was whole- tensor), out-of-range axis -> NnError; fixes densepose per-pixel probs. - translator.rs apply_attention: real scaled-dot-product attention (was a uniform 1/seq_len stub that made any "with attention" ablation == without); mis-shaped checkpoint projections rejected. Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruv
@@ -4,11 +4,39 @@
|
||||
//! different backends (ONNX, tch, Candle).
|
||||
|
||||
use crate::error::{NnError, NnResult};
|
||||
use ndarray::{Array1, Array2, Array3, Array4, ArrayD};
|
||||
use ndarray::{Array1, Array2, Array3, Array4, ArrayD, ArrayViewMutD, Axis};
|
||||
// num_traits is available if needed for advanced tensor operations
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fmt;
|
||||
|
||||
/// Apply a numerically-stable softmax in place to every 1-D lane of `view`
|
||||
/// taken along `axis`. Each lane is shifted by its own max before
|
||||
/// exponentiation, then divided by its own sum, so every lane sums to 1.0
|
||||
/// independently — the per-pixel / per-class normalization densepose needs.
|
||||
///
|
||||
/// `axis` MUST be validated as in-range by the caller.
|
||||
fn softmax_inplace_along_axis(mut view: ArrayViewMutD<'_, f32>, axis: usize) {
|
||||
for mut lane in view.lanes_mut(Axis(axis)) {
|
||||
let max = lane.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
||||
// An all-`-inf` (or empty) lane has no finite max; leave it untouched
|
||||
// to avoid producing NaNs from `exp(-inf - -inf)`.
|
||||
if !max.is_finite() {
|
||||
continue;
|
||||
}
|
||||
let mut sum = 0.0f32;
|
||||
for v in lane.iter_mut() {
|
||||
let e = (*v - max).exp();
|
||||
*v = e;
|
||||
sum += e;
|
||||
}
|
||||
if sum > 0.0 {
|
||||
for v in lane.iter_mut() {
|
||||
*v /= sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shape of a tensor
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub struct TensorShape(Vec<usize>);
|
||||
@@ -288,14 +316,39 @@ impl Tensor {
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply softmax along axis
|
||||
pub fn softmax(&self, _axis: usize) -> NnResult<Tensor> {
|
||||
/// Apply softmax along the given `axis`.
|
||||
///
|
||||
/// Each 1-D lane along `axis` is normalized independently so it sums to
|
||||
/// 1.0. This is the correct semantics for per-pixel / per-class probability
|
||||
/// maps (e.g. DensePose body-part logits over the channel axis). A
|
||||
/// numerically-stable max-shift is applied per lane.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns [`NnError`] if `axis` is out of range for the tensor's rank, or
|
||||
/// if the tensor type is unsupported.
|
||||
pub fn softmax(&self, axis: usize) -> NnResult<Tensor> {
|
||||
match self {
|
||||
Tensor::Float4D(a) => {
|
||||
let max = a.fold(f32::NEG_INFINITY, |acc, &x| acc.max(x));
|
||||
let exp = a.mapv(|x| (x - max).exp());
|
||||
let sum = exp.sum();
|
||||
Ok(Tensor::Float4D(exp / sum))
|
||||
if axis >= a.ndim() {
|
||||
return Err(NnError::tensor_op(format!(
|
||||
"softmax axis {axis} out of range for {}-D tensor",
|
||||
a.ndim()
|
||||
)));
|
||||
}
|
||||
let mut out = a.clone();
|
||||
softmax_inplace_along_axis(out.view_mut().into_dyn(), axis);
|
||||
Ok(Tensor::Float4D(out))
|
||||
}
|
||||
Tensor::FloatND(a) => {
|
||||
if axis >= a.ndim() {
|
||||
return Err(NnError::tensor_op(format!(
|
||||
"softmax axis {axis} out of range for {}-D tensor",
|
||||
a.ndim()
|
||||
)));
|
||||
}
|
||||
let mut out = a.clone();
|
||||
softmax_inplace_along_axis(out.view_mut(), axis);
|
||||
Ok(Tensor::FloatND(out))
|
||||
}
|
||||
_ => Err(NnError::tensor_op(
|
||||
"Softmax not supported for this tensor type",
|
||||
@@ -517,6 +570,67 @@ mod tests {
|
||||
assert!(sigmoid.max().unwrap() < 1.0);
|
||||
}
|
||||
|
||||
// ADR-155 §Tier-2: softmax(axis) must normalize along the GIVEN axis
|
||||
// (per-lane sum == 1), not over the whole tensor.
|
||||
#[test]
|
||||
fn test_softmax_axis_sums_to_one_per_lane() {
|
||||
// 2x3x1x1 tensor; softmax along axis 1 (the size-3 axis).
|
||||
let arr =
|
||||
Array4::from_shape_vec([2, 3, 1, 1], vec![1.0f32, 2.0, 3.0, -1.0, 0.0, 1.0]).unwrap();
|
||||
let t = Tensor::Float4D(arr);
|
||||
let sm = t.softmax(1).unwrap();
|
||||
let out = sm.as_array4().unwrap();
|
||||
// Each lane along axis 1 must sum to 1.0.
|
||||
for b in 0..2 {
|
||||
let lane_sum: f32 = (0..3).map(|c| out[[b, c, 0, 0]]).sum();
|
||||
assert!((lane_sum - 1.0).abs() < 1e-6, "lane {b} sum = {lane_sum}");
|
||||
}
|
||||
// Probabilities must be ordered like the logits within a lane.
|
||||
assert!(out[[0, 0, 0, 0]] < out[[0, 1, 0, 0]]);
|
||||
assert!(out[[0, 1, 0, 0]] < out[[0, 2, 0, 0]]);
|
||||
}
|
||||
|
||||
// ADR-155 §Tier-2: softmax along different axes must give different
|
||||
// results — the old global-softmax bug ignored the axis entirely.
|
||||
#[test]
|
||||
fn test_softmax_axis_choice_matters() {
|
||||
let arr = Array4::from_shape_vec([1, 2, 2, 1], vec![1.0f32, 2.0, 3.0, 4.0]).unwrap();
|
||||
let t = Tensor::Float4D(arr);
|
||||
let along1 = t.softmax(1).unwrap();
|
||||
let along2 = t.softmax(2).unwrap();
|
||||
let a1 = along1.as_array4().unwrap();
|
||||
let a2 = along2.as_array4().unwrap();
|
||||
// The two normalizations partition the values differently, so at least
|
||||
// one element must differ.
|
||||
let mut differs = false;
|
||||
for h in 0..2 {
|
||||
if (a1[[0, 0, h, 0]] - a2[[0, 0, h, 0]]).abs() > 1e-6 {
|
||||
differs = true;
|
||||
}
|
||||
}
|
||||
assert!(differs, "softmax along axis 1 must differ from axis 2");
|
||||
}
|
||||
|
||||
// ADR-155 §Tier-2: known-value check on a tiny tensor.
|
||||
#[test]
|
||||
fn test_softmax_known_values() {
|
||||
// Lane [0, ln(3)] along axis 1 → softmax = [1/4, 3/4].
|
||||
let arr = Array4::from_shape_vec([1, 2, 1, 1], vec![0.0f32, 3.0f32.ln()]).unwrap();
|
||||
let t = Tensor::Float4D(arr);
|
||||
let out = t.softmax(1).unwrap();
|
||||
let a = out.as_array4().unwrap();
|
||||
assert!((a[[0, 0, 0, 0]] - 0.25).abs() < 1e-6);
|
||||
assert!((a[[0, 1, 0, 0]] - 0.75).abs() < 1e-6);
|
||||
}
|
||||
|
||||
// ADR-155 §Tier-2: out-of-range axis must return an error, never panic.
|
||||
#[test]
|
||||
fn test_softmax_axis_out_of_range_errors() {
|
||||
let t = Tensor::zeros_4d([1, 2, 2, 2]);
|
||||
assert!(t.softmax(4).is_err());
|
||||
assert!(t.softmax(99).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_broadcast_compatible() {
|
||||
let a = TensorShape::new(vec![1, 3, 224, 224]);
|
||||
|
||||
@@ -556,34 +556,122 @@ impl ModalityTranslator {
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply multi-head attention
|
||||
/// Apply single-head scaled-dot-product attention over the spatial
|
||||
/// sequence: `softmax(Q·Kᵀ / √d) · V`, with `Q/K/V` linear projections of
|
||||
/// each token's channel vector and a final output projection.
|
||||
///
|
||||
/// The spatial grid `[B, C, H, W]` is treated as a length-`H·W` token
|
||||
/// sequence of `C`-dim feature vectors. Each `*_weight` projection is a
|
||||
/// `[C × C]` matrix applied per token. This is a genuine attention
|
||||
/// operation (not the previous uniform-weight identity stub), so the
|
||||
/// returned per-pair attention weights actually depend on the input.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if any projection weight is not `[C × C]`, so a
|
||||
/// mis-shaped checkpoint can never be silently treated as a no-op.
|
||||
fn apply_attention(
|
||||
&self,
|
||||
input: &Array4<f32>,
|
||||
_weights: &AttentionWeights,
|
||||
weights: &AttentionWeights,
|
||||
) -> NnResult<(Array4<f32>, Array4<f32>)> {
|
||||
let (batch, channels, height, width) = input.dim();
|
||||
let seq_len = height * width;
|
||||
|
||||
// Flatten spatial dimensions
|
||||
let mut flat = ndarray::Array2::zeros((batch, seq_len * channels));
|
||||
// Every projection must be a square [C × C] matrix to act per token.
|
||||
for (name, w) in [
|
||||
("query_weight", &weights.query_weight),
|
||||
("key_weight", &weights.key_weight),
|
||||
("value_weight", &weights.value_weight),
|
||||
("output_weight", &weights.output_weight),
|
||||
] {
|
||||
if w.dim() != (channels, channels) {
|
||||
return Err(NnError::invalid_input(format!(
|
||||
"attention {name} must be [{channels} x {channels}], got [{} x {}]",
|
||||
w.dim().0,
|
||||
w.dim().1
|
||||
)));
|
||||
}
|
||||
}
|
||||
if weights.output_bias.len() != channels {
|
||||
return Err(NnError::shape_mismatch(
|
||||
vec![channels],
|
||||
vec![weights.output_bias.len()],
|
||||
));
|
||||
}
|
||||
|
||||
// Flatten spatial grid into a [seq_len, channels] token matrix per batch.
|
||||
// Project to Q, K, V; compute scaled-dot-product attention; project out.
|
||||
let scale = 1.0 / (channels as f32).sqrt();
|
||||
let mut out = Array4::zeros((batch, channels, height, width));
|
||||
let mut attention_weights = Array4::zeros((batch, 1, seq_len, seq_len));
|
||||
|
||||
for b in 0..batch {
|
||||
// Tokens: [seq_len, channels].
|
||||
let mut tokens = ndarray::Array2::<f32>::zeros((seq_len, channels));
|
||||
for h in 0..height {
|
||||
for w in 0..width {
|
||||
let s = h * width + w;
|
||||
for c in 0..channels {
|
||||
flat[[b, (h * width + w) * channels + c]] = input[[b, c, h, w]];
|
||||
tokens[[s, c]] = input[[b, c, h, w]];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Q = tokens·Wqᵀ, etc. (row vector × [C×C] projection).
|
||||
let q = tokens.dot(&weights.query_weight.t());
|
||||
let k = tokens.dot(&weights.key_weight.t());
|
||||
let v = tokens.dot(&weights.value_weight.t());
|
||||
|
||||
// Scores = softmax_row(Q·Kᵀ · scale), then context = Scores·V.
|
||||
let scores = q.dot(&k.t()).mapv(|x| x * scale);
|
||||
for i in 0..seq_len {
|
||||
// Numerically-stable row softmax.
|
||||
let mut max = f32::NEG_INFINITY;
|
||||
for j in 0..seq_len {
|
||||
max = max.max(scores[[i, j]]);
|
||||
}
|
||||
let mut sum = 0.0f32;
|
||||
let mut row = vec![0.0f32; seq_len];
|
||||
for j in 0..seq_len {
|
||||
let e = (scores[[i, j]] - max).exp();
|
||||
row[j] = e;
|
||||
sum += e;
|
||||
}
|
||||
if sum > 0.0 {
|
||||
for j in 0..seq_len {
|
||||
row[j] /= sum;
|
||||
}
|
||||
}
|
||||
for j in 0..seq_len {
|
||||
attention_weights[[b, 0, i, j]] = row[j];
|
||||
}
|
||||
}
|
||||
|
||||
// Context = attention · V, then output projection + bias.
|
||||
for h in 0..height {
|
||||
for w in 0..width {
|
||||
let i = h * width + w;
|
||||
// ctx[c] = Σ_j attn[i,j] · v[j,c]
|
||||
let mut ctx = vec![0.0f32; channels];
|
||||
for j in 0..seq_len {
|
||||
let a = attention_weights[[b, 0, i, j]];
|
||||
for c in 0..channels {
|
||||
ctx[c] += a * v[[j, c]];
|
||||
}
|
||||
}
|
||||
// out[c] = Σ_c' ctx[c'] · Wo[c, c'] + bias[c]
|
||||
for c in 0..channels {
|
||||
let mut acc = weights.output_bias[c];
|
||||
for cp in 0..channels {
|
||||
acc += ctx[cp] * weights.output_weight[[c, cp]];
|
||||
}
|
||||
out[[b, c, h, w]] = acc;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// For simplicity, return input unchanged with identity attention
|
||||
let attention_weights = Array4::from_elem(
|
||||
(batch, self.config.attention_heads, seq_len, seq_len),
|
||||
1.0 / seq_len as f32,
|
||||
);
|
||||
|
||||
Ok((input.clone(), attention_weights))
|
||||
Ok((out, attention_weights))
|
||||
}
|
||||
|
||||
/// Compute translation loss between predicted and target features
|
||||
@@ -760,6 +848,76 @@ mod tests {
|
||||
assert_eq!(config.activation, ActivationType::GELU);
|
||||
}
|
||||
|
||||
// ADR-155 §Tier-2: apply_attention must perform real scaled-dot-product
|
||||
// attention, not return uniform 1/seq_len weights. With identity Q/K/V
|
||||
// projections and a non-uniform input, the attention weights must NOT all
|
||||
// equal 1/seq_len, and each row must still be a valid distribution.
|
||||
#[test]
|
||||
fn test_attention_is_not_uniform_stub() {
|
||||
let channels = 4usize;
|
||||
let height = 2usize;
|
||||
let width = 2usize;
|
||||
let seq_len = height * width;
|
||||
|
||||
// Identity projections so Q=K=V=tokens; output = identity, zero bias.
|
||||
let identity = ndarray::Array2::<f32>::eye(channels);
|
||||
let weights = AttentionWeights {
|
||||
query_weight: identity.clone(),
|
||||
key_weight: identity.clone(),
|
||||
value_weight: identity.clone(),
|
||||
output_weight: identity,
|
||||
output_bias: ndarray::Array1::zeros(channels),
|
||||
};
|
||||
|
||||
// Non-uniform input: each spatial location has a distinct feature vector.
|
||||
let mut input = Array4::<f32>::zeros((1, channels, height, width));
|
||||
for c in 0..channels {
|
||||
for h in 0..height {
|
||||
for w in 0..width {
|
||||
input[[0, c, h, w]] = (c + 2 * h + 4 * w) as f32;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let config = TranslatorConfig::default().with_attention(1);
|
||||
let translator = ModalityTranslator::new(config).unwrap();
|
||||
let (out, attn) = translator.apply_attention(&input, &weights).unwrap();
|
||||
|
||||
// Each attention row must sum to 1 (valid softmax distribution).
|
||||
for i in 0..seq_len {
|
||||
let row_sum: f32 = (0..seq_len).map(|j| attn[[0, 0, i, j]]).sum();
|
||||
assert!((row_sum - 1.0).abs() < 1e-5, "row {i} sum = {row_sum}");
|
||||
}
|
||||
// Weights must NOT all be the uniform 1/seq_len value of the old stub.
|
||||
let uniform = 1.0 / seq_len as f32;
|
||||
let any_non_uniform = (0..seq_len)
|
||||
.flat_map(|i| (0..seq_len).map(move |j| (i, j)))
|
||||
.any(|(i, j)| (attn[[0, 0, i, j]] - uniform).abs() > 1e-4);
|
||||
assert!(any_non_uniform, "attention collapsed to uniform stub");
|
||||
// Output is finite and shaped like the input.
|
||||
assert_eq!(out.dim(), input.dim());
|
||||
assert!(out.iter().all(|v| v.is_finite()));
|
||||
}
|
||||
|
||||
// ADR-155 §Tier-2: a mis-shaped projection weight must be rejected, never
|
||||
// silently treated as a no-op.
|
||||
#[test]
|
||||
fn test_attention_rejects_wrong_weight_shape() {
|
||||
let channels = 4usize;
|
||||
let bad = ndarray::Array2::<f32>::zeros((channels + 1, channels));
|
||||
let weights = AttentionWeights {
|
||||
query_weight: bad.clone(),
|
||||
key_weight: bad.clone(),
|
||||
value_weight: bad.clone(),
|
||||
output_weight: bad,
|
||||
output_bias: ndarray::Array1::zeros(channels),
|
||||
};
|
||||
let input = Array4::<f32>::zeros((1, channels, 2, 2));
|
||||
let config = TranslatorConfig::default().with_attention(1);
|
||||
let translator = ModalityTranslator::new(config).unwrap();
|
||||
assert!(translator.apply_attention(&input, &weights).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_loss_computation() {
|
||||
let config = TranslatorConfig::default();
|
||||
|
||||
Reference in New Issue
Block a user