Hello PyTorch Community and Team,
I want to share a mathematical framework and custom layer extension I’ve been developing to address the bottlenecks of static training loops and vanishing gradients: the Tiwantic-Feedback Equation (TFE).
Instead of relying solely on external learning rate schedulers to manage optimization pacing, TFE embeds an autonomous feedback mechanism directly into the raw forward and backward paths of the layer itself.
The Mathematical Concept
The core equation is formulated as:
$$f(x) = x \cdot \sigma(w \cdot x) + \alpha \cdot b^2$$
Where:
-
$\sigma(w \cdot x)$ acts as an input-driven gating function.
-
$\alpha$ is a dynamic feedback coefficient tied to the model’s real-time training loss variance.
-
$b^2$ serves as a baseline stabilization factor.
How It Works Under the Hood
-
High Loss Variance (Acceleration State): When the model is initialized or struggling with high training errors, $\alpha$ scales up. This amplifies the gradient flow during the backward pass, helping the weights forcefully break out of bad local minima or random states.
-
Low Loss Variance (High-Precision State): As the loss drops and stabilizes, $\alpha$ smoothly trends down. The layer naturally dampens its gradient magnitudes, allowing for high-precision micro-adjustments without overshooting the global minimum.
-
Memory Shielding: Decoupling input processing from the baseline shift ($\alpha \cdot b^2$) creates a buffer that helps protect older weight patterns during continuous learning loops.
PyTorch Autograd Native Implementation
I have implemented this natively using torch.autograd.Function so it integrates seamlessly into PyTorch’s computational graph without breaking custom autograd pipelines:
Python
import torch
import torch.nn as nn
class TFEScaleFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, x, weight, bias, alpha):
ctx.save_for_backward(x, weight, bias, alpha)
# Core TFE Forward calculation
output = x * torch.sigmoid(weight * x) + alpha * (bias ** 2)
return output
@staticmethod
def backward(ctx, grad_output):
x, weight, bias, alpha = ctx.saved_tensors
sig = torch.sigmoid(weight * x)
dsig = sig * (1 - sig)
# Explicit analytical derivatives for the backward pass
grad_x = grad_output * (sig + x * dsig * weight)
grad_weight = grad_output * (x * x * dsig)
grad_bias = grad_output * (2 * alpha * bias)
grad_alpha = grad_output * (bias ** 2)
return grad_x, grad_weight, grad_bias, grad_alpha
class TiwanticLayer(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.randn(1, requires_grad=True))
self.bias = nn.Parameter(torch.randn(1, requires_grad=True))
def forward(self, x, alpha):
return TFEScaleFunction.apply(x, self.weight, self.bias, alpha)
What I’m Looking For From the Community:
I am aiming to package this into an open-source extension (tiwantic-core) and would love to get feedback from the PyTorch team and community on a couple of points:
-
Autograd Optimization: Are there any vectorization or memory-efficient tensor operations I can apply to the backward pass derivatives to make this run even faster on CUDA?
-
Benchmarking Collaboration: I am looking to test this inside deeper architecture blocks (like Transformer multi-head attention inputs) to see how effectively it stabilizes learning curves compared to standard LayerNorm + ReLU/GELU steps.