Nimai: Controlled Augmentation for Security
- Nimai is a generative data augmentation framework for ML security classifiers that addresses data scarcity, imbalance, and concept drift in tabular data.
- It leverages a VQ-VAE combined with a masked-token model to achieve controlled synthesis through sample-conditioned and class-conditioned modes.
- Controlled augmentation with Nimai has demonstrated up to a 32.61% improvement in constrained settings and 60.4% post-deployment recovery under drift.
Searching arXiv for the specified paper to ground the article in the cited source. Nimai is a generative data augmentation framework tailored for ML-based security classifiers on tabular data. It was introduced to provide highly controlled synthesis of minority and underrepresented patterns so that downstream classifiers generalize better under real-world data challenges typical in security, including imbalanced classes, scarce or biased attack data, high-dimensional features, and concept drift. The framework combines a VQ-VAE with a masked-token modeling prior to support both sample-conditioned and class-conditioned generation, and the reported study evaluates it across seven security tasks using six state-of-the-art GenAI baselines. In the reported experiments, controlled augmentation with Nimai achieved improvements of up to in severely data-constrained settings and up to post-deployment recovery with minimal labels under drift (Kanchi et al., 8 Jul 2025).
1. Problem Setting and Motivation
Nimai was motivated by a diagnosis that the main bottlenecks in many security classifiers are data-centric rather than purely algorithmic. In a systematic study of 35 security papers, 32 reported data challenges. The most common were limited or biased training data, class imbalance, underrepresented attack patterns, few labeled samples, high dimensionality, and concept drift. The paper’s central argument is that better classifiers cannot recover information that is absent, biased, or severely imbalanced in the training distribution; without addressing the data, gains saturate (Kanchi et al., 8 Jul 2025).
The security-specific challenges are described in operational terms. Class imbalance arises because rare attacks are contrasted with abundant benign data, making minority classes hard to learn. Underrepresented attack patterns create holes within a class distribution and degrade recall on rare behaviors. Few labeled samples reflect costly expert labeling and feature extraction. High-dimensional features increase sample complexity. Concept drift reflects the fact that attackers adapt and the test distribution diverges from training, requiring rapid post-deployment recovery. The same study also identifies label noise, overlapping class distributions, and sparse feature vectors as impediments under which GenAI itself may struggle.
Within this framing, generative augmentation is assigned a specific role: learning the empirical distribution and synthesizing additional training points that address scarcity and imbalance. The paper distinguishes between targeting rare sub-regions of a class distribution and generating broader class-consistent diversity. Existing off-the-shelf tabular GenAI models are described as lacking fine-grained control to target specific regions of the data manifold. Nimai was created to fill that gap with sample-conditioned generation while also supporting class-conditioned generation.
2. Controlled Synthesis Paradigm
The defining feature of Nimai is control over where synthetic data are generated. “Control” is operationalized through sample-conditioning, class-conditioning, class rebalancing, and a tunable synthesis mode governed by a masking ratio in a discrete latent space (Kanchi et al., 8 Jul 2025).
| Mode | Conditioning scheme | Intended effect |
|---|---|---|
| Nimai-S | Condition on a specific real sample and class label | Strengthen underrepresented regions within a class |
| Nimai-C | Condition only on the class label | Improve balance and diversity; explore shifted regions under drift |
| Nimai-hybrid | Use Nimai-S on a few newly labeled uncertain points, then Nimai-C for the remainder | Rapid adaptation to concept drift with minimal labeling |
In Nimai-S, a real sample is encoded into discrete latent tokens, a subset of those tokens is masked, and the masked-token model fills the missing tokens conditioned on the remaining tokens and the class label. The resulting synthetic example stays near the anchor sample, and the synthesized point preserves the label of the conditioned sample. This mode is explicitly designed for minority and underrepresented patterns.
In Nimai-C, all latent tokens are masked and the masked-token model fills the entire latent sequence conditioned only on the class label. This allows the generator to explore broader class-consistent regions rather than staying local to an observed point. The paper presents this as particularly relevant when diversity within a class, rather than local repair around existing samples, is the dominant requirement.
The masking ratio serves as a control knob that interpolates between the two modes. At one extreme, masking all tokens recovers class-conditioning; at partial masking, the synthesis remains partially anchored to a real sample. Class balancing is then implemented by generating synthetic data until each minority class matches the size of the largest class.
3. Architecture, Training Stages, and Objectives
Nimai’s high-level design is a VAE with a discrete latent space, implemented as a VQ-VAE, coupled to a masked-token modeling prior. The stated rationale is that discrete latents reduce posterior collapse and offer precise control similar to LLMs over tokens, but for tabular codes. The encoder and decoder are transformer blocks, and tabular features are scaled to for stable learning (Kanchi et al., 8 Jul 2025).
The discrete latent space is defined by a learnable codebook that maps encoder outputs to nearest code vectors, with the corresponding indices functioning as discrete tokens. A BERT-style masked-token model is then trained over the quantized latent sequences together with class labels. Because the prior is non-autoregressive, masked tokens can be predicted in a few steps rather than through sequential generation.
Training is staged. First, the VQ-VAE components—encoder, codebook, and decoder—are trained on real data to minimize reconstruction and codebook losses. Second, the masked-token prior is trained on the quantized latent sequences with class labels using a cross-entropy objective for masked-token prediction. Third, synthesis proceeds by masking all or part of the latent sequence, sampling replacements from the prior, and decoding back to tabular space.
The paper reports the following objectives:
$\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$
$\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$
$\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$
$\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$
$\mathcal{L}_{\text{MTM} = \mathbb{E}_{(\mathbf{y}, c) \sim D} \left[ - \sum_{i} \log p\left( \mathbf{y}_i \mid \mathbf{Y}_{\text{mask}, c \right) \right]$
For robustness under scarce tabular security data, the reported implementation uses early stopping on validation loss and ASHA hyperparameter search over both the VQ-VAE and the masked-token model. The search excludes settings that cause codebook collapse. The paper also states that discrete latents are less prone to posterior collapse and can behave more stably with scarce data than continuous-latent VAEs.
4. Evaluation Corpus and Protocol
The experimental study spans seven security tasks with tabular features and either binary or multiclass labels. These tasks were selected to cover scarcity, imbalance, drift, sparsity, and overlap (Kanchi et al., 8 Jul 2025).
| Task | Setting | Noted challenge |
|---|---|---|
| BGP serial hijacker detection | Binary; 52 features; 180 training samples with only 17 malicious | Extreme scarcity; malicious F-score emphasized |
| Tor-based malware vs benign traffic | Binary; 175 features; 8,115 train | Overlapping class distributions |
| IoT cryptojacking detection | Binary; 273 features; 43,866 train with 1,116 attack | Strong class imbalance |
| Cookie privacy classification | Multiclass, 4 classes; 1,689 features | Noisy labels and high sparsity |
| MS-Malware | Multiclass, 9 classes; 1,804 features; training reduced to 434 | Scarcity plus high dimensionality |
| BODMAS | Multiclass, 40 classes; 2,831 features; train on month 0, test on months 1–12 | Concept drift |
| nPrintML OS fingerprinting | Multiclass, 13 classes; 4,169 sparse binary features; 8,271 train | Extreme sparsity |
The evaluation metric is macro F-score except for BGP and Tor, where malicious-class F-score is reported. The paper gives the macro F-score as
$F = \dfrac{2}{N}\sum_{i=1}^{N}\frac{ P_{i} \cdot R_i}{P_{i} + R_{i}$
and the relative gain over training on real data alone as
$\Delta G = \frac{F_{aug} - F_{real} {F_{real}$
All results are reported over 10 random seeds with mean and standard deviation. A result is deemed “unreliable” if the coefficient of variation is at least 1. Statistical ranking uses Kruskal–Wallis and Dunn’s tests as needed.
The baselines comprise six state-of-the-art GenAI methods and two traditional oversampling comparators. The GenAI baselines are TVAE, CTAB-GAN+, TabDDPM, TabSyn, GReaT, and REaLTabFormer. The traditional comparators are SMOTE and MC-CCR. Their descriptions in the study are method-specific: TVAE is a VAE for tabular data using Variational Gaussian Mixtures for multimodal continuous features; CTAB-GAN+ is a conditional GAN with Wasserstein loss and gradient penalty; TabDDPM is a diffusion model with a denoising MLP and TPE hyperparameter search; TabSyn combines a VAE with diffusion in latent space; GReaT uses DistilGPT-2 over a text serialization of tables; and REaLTabFormer is an LLM-based tabular generator with sub-columnar tokens, early stopping using similarity metrics, and token-constrained decoding.
5. Empirical Performance and Mechanistic Interpretation
The strongest pre-deployment gains are task-dependent and align with the intended control modes. On IoT, a binary task with severe imbalance, Nimai-S achieved the best mean 0, outperforming SMOTE at 1, CTAB-GAN+ at 2, Nimai-C at 3, and TVAE at 4. TabDDPM failed through mode collapse, and GReaT and REaLTabFormer were infeasible or failed to converge in time. On BGP, where the training set contains only about 180 samples and 17 malicious examples, Nimai-S achieved 5, Nimai-C 6, CTAB-GAN+ 7, and SMOTE 8, while TVAE, TabSyn, GReaT, and REaLTabFormer were negative and TabDDPM was near zero (Kanchi et al., 8 Jul 2025).
On MS-Malware, a high-dimensional multiclass task with the training set reduced to 434, class-conditioned generators were strongest: TabDDPM achieved 9, CTAB-GAN+ 0, Nimai-C 1, TVAE 2, and Nimai-S 3, while SMOTE was slightly negative at 4. On BODMAS, which exhibits concept drift across 12 test months, a GenAI method improved performance in 8 of 12 months. Nimai-C was best in 6 of the 8 months with positive gains, including month 5 at 5 and month 9 at 6, and it was also positive in months 2, 3, 10, and 11 with gains of 7, 8, 9, and $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$0, respectively. Month 6 was led by TVAE at $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$1, with Nimai-C also positive at $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$2.
The tasks on which augmentation failed are equally important. On Tor, only REaLTabFormer improved reliably, at $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$3, while most methods degraded, suggesting that overlap or noise undermined augmentation benefits. On Cookie, which is reported as $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$4 sparse and contains noisy labels, and on nPrintML, which is reported as $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$5 sparse, no reliable GenAI gains were observed; SMOTE gave a small positive gain on Cookie at $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$6. The study notes evidence of noisy or contradictory labels in Cookie, including 3,707 identical feature vectors mapped to multiple labels, and poor separability across OS variants in nPrintML. A controlled stress test on BGP in which 20% label flips were induced drove all augmentation methods to negative gains, including Nimai-S at $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$7, CTAB-GAN+ at $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$8, and SMOTE at $\mathcal{L}_{\text{recon} = \| x - \hat{x} \|^2$9.
The paper also provides a mechanistic analysis of why augmentation helped on IoT and BGP. Using LIME on 50 corrected test instances per task, it reports that for influential features, IoT showed a $\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$0 increase in entropy and a $\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$1 absolute reduction in skewness toward 0, while BGP showed a $\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$2 increase in entropy and a $\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$3 absolute skewness reduction. The authors interpret these results as indicating that more symmetric, higher-entropy feature distributions improve generalization.
A separate empirical finding concerns trainability. Five of the six existing GenAI models failed on at least one task. The LLM-based generators GReaT and REaLTabFormer encountered token limits and memory blowups on high-dimensional tables, where more than 1,000 columns often become more than 10,000 tokens; on IoT, GReaT was estimated at 5,714 GPU-hours per epoch. TabDDPM exhibited mode collapse on IoT, TabSyn ran out of memory for very high dimensions or class counts, and TVAE’s rejection sampling could not produce enough class-conditional samples for Cookie even after 100 hours.
6. Post-Deployment Adaptation Under Concept Drift
The study extends Nimai to a post-deployment setting through Nimai-hybrid, with the objective of rapid adaptation using minimal labeling. The proposed workflow begins by detecting drift with methods such as CADE or Transcend. On a drifted month $\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$4, uncertain test samples are selected by confidence, for example in the range $\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$5–$\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$6, from a small subset of about $\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$7. Only a few of these samples—reported as 9 to 64—are newly labeled (Kanchi et al., 8 Jul 2025).
Nimai-hybrid then combines both control modes. Nimai-S is used to generate $\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$8 synthetic samples anchored to the few newly labeled uncertain points so as to capture the new modes. Nimai-C is then used to generate the remainder needed to rebalance classes and preserve wider class diversity. The classifier is retrained on the union of real and hybrid synthetic data.
The evaluation focuses on the worst drift months in BODMAS, months 5 and 6. Without GenAI, the macro F-score was 36.40 and 33.60. INSOMNIA, which retrains on real uncertainty samples only, achieved 34.40 and 33.40. Nimai-C trained from month 0 achieved 58.03 $\mathcal{L}_{\text{commit} = \beta \| z_e(x) - \text{sg}(e) \|^2$9 and 43.09 $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$0. Nimai-hybrid achieved 59.78 $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$1 on month 5 with 64 labeled samples, and 53.92 $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$2 on month 6 with as few as 9–64 labeled samples depending on which prior month provided uncertainty seeds. The paper summarizes this as up to a $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$3 performance boost.
The reported ablation is consistent with the pre-deployment results: sample-conditioning excels in extreme class imbalance and scarcity, class-conditioning excels under concept drift and high-dimensional multiclass settings, and combining them yields fast recovery when only a small number of new labels are available.
7. Limitations, Deployment Conditions, and Future Directions
The reported limitations are explicit. Nimai’s gains are limited when labels are noisy, class distributions heavily overlap, or feature vectors are extremely sparse, as observed in Cookie and nPrintML. The framework currently operates in feature space only, not problem space, and mixed data-type encodings are standardized rather than type-specific (Kanchi et al., 8 Jul 2025).
The deployment guidance in the study is regime-specific. Nimai-S is recommended when severe class imbalance and underrepresented attack patterns are the dominant issue. Nimai-C is recommended for high-dimensional multiclass settings with limited data and for concept drift scenarios in which broader class diversity helps. Nimai-hybrid is recommended for quick post-deployment adaptation from a handful of new labels. Conversely, significant label noise, heavy class overlap, or extremely sparse feature vectors with weak class separability are presented as cases in which improving feature engineering, label cleaning, or problem framing may be more effective than augmentation.
Operationally, the preprocessing requirement is min–max scaling of all features to $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$4. The downstream models are described as robust to this choice, and Logistic Regression on IoT is specifically reported to improve markedly after scaling. The implementation uses a transformer encoder and decoder, a VQ-VAE with a codebook, a BERT-style masked-token prior, early stopping on validation loss, ASHA for hyperparameter search, and codebook-collapse checks during search. The hyperparameter search space includes attention heads $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$5, FFN size $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$6, transformer layers $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$7, codebook size $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$8, latent length $\mathcal{L}_{\text{embed} = \| \text{sg}(z_e(x)) - e \|^2$9 with $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$0 in $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$1, embedding size $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$2, loss weights $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$3, and EMA decay $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$4. The masking ratio is the main synthesis hyperparameter because it controls interpolation between Nimai-S and Nimai-C, while codebook size, embedding dimension, and the reconstruction and codebook loss weights $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$5 mediate the trade-off between diversity, fidelity, and stability.
The reported compute profile is comparatively modest on an A100: average training time per epoch is about $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$6 seconds, with per-task examples of about $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$7 seconds for BGP, $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$8 seconds for MS-Malware, $\mathcal{L} = \mathcal{L}_{\text{recon} + \alpha \left( \mathcal{L}_{\text{embed} + \mathcal{L}_{\text{commit} \right)$9 second for BODMAS, $\mathcal{L}_{\text{MTM} = \mathbb{E}_{(\mathbf{y}, c) \sim D} \left[ - \sum_{i} \log p\left( \mathbf{y}_i \mid \mathbf{Y}_{\text{mask}, c \right) \right]$0 seconds for Tor, $\mathcal{L}_{\text{MTM} = \mathbb{E}_{(\mathbf{y}, c) \sim D} \left[ - \sum_{i} \log p\left( \mathbf{y}_i \mid \mathbf{Y}_{\text{mask}, c \right) \right]$1 seconds for Cookie, $\mathcal{L}_{\text{MTM} = \mathbb{E}_{(\mathbf{y}, c) \sim D} \left[ - \sum_{i} \log p\left( \mathbf{y}_i \mid \mathbf{Y}_{\text{mask}, c \right) \right]$2 seconds for IoT, and $\mathcal{L}_{\text{MTM} = \mathbb{E}_{(\mathbf{y}, c) \sim D} \left[ - \sum_{i} \log p\left( \mathbf{y}_i \mid \mathbf{Y}_{\text{mask}, c \right) \right]$3 seconds for nPrintML. Risk considerations in the paper include synthetic bias and artifact overfitting, the need to evaluate on a held-out real test set and use statistical tests for robust gains, privacy and dual-use concerns despite the focus on feature-space synthesis for defenders, and the possibility of drift mis-diagnosis if uncertain-sample selection is incorrect.
The paper identifies three forward directions: scaling LLM-based tabular generators through long-context and memory-efficient designs, advanced diffusion, VAE, and GAN strategies for rare-sample synthesis through strategic sampling and weighted learning, and improved handling of mixed data types, label noise, overlap, and adversarially robust augmentation. It also states that all code and datasets for Nimai and the evaluated classifiers will be released.